Qwen 3, Bug Fixes (#2445)
* bug fix #2008 (#2039) * fix (#2051) * Update loader.py * Update pyproject.toml * Update pyproject.toml * Update vision.py * more prints * Update loader.py * LoRA 16bit fix * Update vision.py * Update vision.py * Update _utils.py * Update vision.py * move forced float32 * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * move print * Update _utils.py * disable bfloat16 * Fix forced float32 * move float32 * Ensure trust_remote_code propegates down to unsloth_compile_transformers (#2075) * Update _utils.py * Show both `peft_error` and `autoconfig_error`, not just `autoconfig_error` (#2080) When loading a PEFT model fails, only the `autoconfig_error` is shown. Instead of the `peft_error`, which is what really matters when we're trying to load a PEFT adapter, the user will see something like this: ``` RuntimeError: Unrecognized model in my_model. Should have a `model_type` key in its config.json, or contain one of the following strings in its name: albert, align, altclip, ... ``` This PR just changes it so `autoconfig_error` and `peft_error` are both displayed. * fix error message (#2046) * Update vision.py * Update _utils.py * Update pyproject.toml * Update __init__.py * Update __init__.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update rl_replacements.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update rl_replacements.py * Update vision.py * Update rl_replacements.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Remove double generate patch * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update mapper.py * Update vision.py * fix: config.torch_dtype in LlamaModel_fast_forward_inference (#2091) * fix: config.torch_dtype in LlamaModel_fast_forward_inference * Update llama.py * update for consistency --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> * versioning * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * model_type_arch * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update loader.py * check * Update _utils.py * Update loader.py * Update loader.py * Remove prints * Update README.md typo * Update _utils.py * Update _utils.py * versioning * Update _utils.py * Update _utils.py * Update _utils.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update vision.py * HF Transfer * fix(utils): add missing importlib import to fix NameError (#2134) This commit fixes a NameError that occurs when `importlib` is referenced in _utils.py without being imported, especially when UNSLOTH_USE_MODELSCOPE=1 is enabled. By adding the missing import statement, the code will no longer throw a NameError. * Add QLoRA Train and Merge16bit Test (#2130) * add reference and unsloth lora merging tests * add test / dataset printing to test scripts * allow running tests from repo root * add qlora test readme * more readme edits * ruff formatting * additional readme comments * forgot to add actual tests * add apache license * Update pyproject.toml * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update loader.py * Update loader.py * Revert * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Update vision.py * Bug fix * Update mapper.py * check SDPA for Mistral 3, Pixtral * Update vision.py * Versioning * Update rl_replacements.py * Update README.md * add model registry * move hf hub utils to unsloth/utils * refactor global model info dicts to dataclasses * fix dataclass init * fix llama registration * remove deprecated key function * start registry reog * add llama vision * quant types -> Enum * remap literal quant types to QuantType Enum * add llama model registration * fix quant tag mapping * add qwen2.5 models to registry * add option to include original model in registry * handle quant types per model size * separate registration of base and instruct llama3.2 * add QwenQVQ to registry * add gemma3 to registry * add phi * add deepseek v3 * add deepseek r1 base * add deepseek r1 zero * add deepseek distill llama * add deepseek distill models * remove redundant code when constructing model names * add mistral small to registry * rename model registration methods * rename deepseek registration methods * refactor naming for mistral and phi * add global register models * refactor model registration tests for new registry apis * add model search method * remove deprecated registration api * add quant type test * add registry readme * make llama registration more specific * clear registry when executing individual model registration file * more registry readme updates * Update _auto_install.py * Llama4 * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Synthetic data * Update mapper.py * Xet and Synthetic * Update synthetic.py * Update loader.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update synthetic.py * Update pyproject.toml * Delete .gitignore --------- Co-authored-by: Mukkesh Ganesh <mukmckenzie@gmail.com> Co-authored-by: Kareem <81531392+KareemMusleh@users.noreply.github.com> Co-authored-by: Xander Hawthorne <167850078+CuppaXanax@users.noreply.github.com> Co-authored-by: Isaac Breen <isaac.breen@icloud.com> Co-authored-by: lurf21 <93976703+lurf21@users.noreply.github.com> Co-authored-by: Jack Shi Wei Lun <87535974+jackswl@users.noreply.github.com> Co-authored-by: naliazheli <nalia0316@gmail.com> Co-authored-by: jeromeku <jerome.ku@gmail.com> Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
599a15ea8d
commit
9a930bb095
19 changed files with 1470 additions and 225 deletions
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
91
tests/test_model_registry.py
Normal file
91
tests/test_model_registry.py
Normal file
|
|
@ -0,0 +1,91 @@
|
||||||
|
"""
|
||||||
|
|
||||||
|
Test model registration methods
|
||||||
|
Checks that model registration methods work for respective models as well as all models
|
||||||
|
The check is performed
|
||||||
|
- by registering the models
|
||||||
|
- checking that the instantiated models can be found on huggingface hub by querying for the model id
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from huggingface_hub import ModelInfo as HfModelInfo
|
||||||
|
|
||||||
|
from unsloth.registry import register_models, search_models
|
||||||
|
from unsloth.registry._deepseek import register_deepseek_models
|
||||||
|
from unsloth.registry._gemma import register_gemma_models
|
||||||
|
from unsloth.registry._llama import register_llama_models
|
||||||
|
from unsloth.registry._mistral import register_mistral_models
|
||||||
|
from unsloth.registry._phi import register_phi_models
|
||||||
|
from unsloth.registry._qwen import register_qwen_models
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, QUANT_TAG_MAP, QuantType
|
||||||
|
from unsloth.utils.hf_hub import get_model_info
|
||||||
|
|
||||||
|
MODEL_NAMES = [
|
||||||
|
"llama",
|
||||||
|
"qwen",
|
||||||
|
"mistral",
|
||||||
|
"phi",
|
||||||
|
"gemma",
|
||||||
|
"deepseek",
|
||||||
|
]
|
||||||
|
MODEL_REGISTRATION_METHODS = [
|
||||||
|
register_llama_models,
|
||||||
|
register_qwen_models,
|
||||||
|
register_mistral_models,
|
||||||
|
register_phi_models,
|
||||||
|
register_gemma_models,
|
||||||
|
register_deepseek_models,
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ModelTestParam:
|
||||||
|
name: str
|
||||||
|
register_models: callable
|
||||||
|
|
||||||
|
|
||||||
|
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, MODEL_REGISTRATION_METHODS)
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
# Test that model registration methods register respective models
|
||||||
|
@pytest.mark.parametrize("model_test_param", TestParams, ids=lambda param: param.name)
|
||||||
|
def test_model_registration(model_test_param: ModelTestParam):
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
registration_method = model_test_param.register_models
|
||||||
|
registration_method()
|
||||||
|
registered_models = MODEL_REGISTRY.keys()
|
||||||
|
missing_models = _test_model_uploaded(registered_models)
|
||||||
|
assert not missing_models, (
|
||||||
|
f"{model_test_param.name} missing following models: {missing_models}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_all_model_registration():
|
||||||
|
register_models()
|
||||||
|
registered_models = MODEL_REGISTRY.keys()
|
||||||
|
missing_models = _test_model_uploaded(registered_models)
|
||||||
|
assert not missing_models, f"Missing following models: {missing_models}"
|
||||||
|
|
||||||
|
def test_quant_type():
|
||||||
|
# Test that the quant_type is correctly set for model paths
|
||||||
|
# NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH
|
||||||
|
dynamic_quant_models = search_models(quant_types=[QuantType.UNSLOTH])
|
||||||
|
assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models)
|
||||||
|
quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH]
|
||||||
|
assert all(quant_tag in m.model_path for m in dynamic_quant_models)
|
||||||
|
|
@ -18,7 +18,7 @@ from packaging.version import Version as V
|
||||||
v = V(torch.__version__)
|
v = V(torch.__version__)
|
||||||
cuda = str(torch.version.cuda)
|
cuda = str(torch.version.cuda)
|
||||||
is_ampere = torch.cuda.get_device_capability()[0] >= 8
|
is_ampere = torch.cuda.get_device_capability()[0] >= 8
|
||||||
if cuda != "12.1" and cuda != "11.8" and cuda != "12.4" and cuda != "12.6": raise RuntimeError(f"CUDA = {cuda} not supported!")
|
if cuda != "12.1" and cuda != "11.8" and cuda != "12.4" and cuda != "12.6" and cuda != "12.8": raise RuntimeError(f"CUDA = {cuda} not supported!")
|
||||||
if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!")
|
if v <= V('2.1.0'): raise RuntimeError(f"Torch = {v} too old!")
|
||||||
elif v <= V('2.1.1'): x = 'cu{}{}-torch211'
|
elif v <= V('2.1.1'): x = 'cu{}{}-torch211'
|
||||||
elif v <= V('2.1.2'): x = 'cu{}{}-torch212'
|
elif v <= V('2.1.2'): x = 'cu{}{}-torch212'
|
||||||
|
|
@ -28,6 +28,7 @@ elif v < V('2.5.0'): x = 'cu{}{}-torch240'
|
||||||
elif v < V('2.5.1'): x = 'cu{}{}-torch250'
|
elif v < V('2.5.1'): x = 'cu{}{}-torch250'
|
||||||
elif v <= V('2.5.1'): x = 'cu{}{}-torch251'
|
elif v <= V('2.5.1'): x = 'cu{}{}-torch251'
|
||||||
elif v < V('2.7.0'): x = 'cu{}{}-torch260'
|
elif v < V('2.7.0'): x = 'cu{}{}-torch260'
|
||||||
|
elif v < V('2.8.0'): x = 'cu{}{}-torch270'
|
||||||
else: raise RuntimeError(f"Torch = {v} too new!")
|
else: raise RuntimeError(f"Torch = {v} too new!")
|
||||||
x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
|
x = x.format(cuda.replace(".", ""), "-ampere" if is_ampere else "")
|
||||||
print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')
|
print(f'pip install --upgrade pip && pip install "unsloth[{x}] @ git+https://github.com/unslothai/unsloth.git"')
|
||||||
|
|
@ -13,10 +13,7 @@
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"check_vllm_status",
|
"SyntheticDataKit",
|
||||||
"async_load_vllm",
|
|
||||||
"destroy_vllm",
|
|
||||||
"configure_synthetic_data_kit",
|
|
||||||
]
|
]
|
||||||
import subprocess
|
import subprocess
|
||||||
import time
|
import time
|
||||||
|
|
@ -26,236 +23,240 @@ import torch
|
||||||
import gc
|
import gc
|
||||||
import time
|
import time
|
||||||
from unsloth_zoo.vllm_utils import load_vllm
|
from unsloth_zoo.vllm_utils import load_vllm
|
||||||
from transformers import AutoConfig
|
from transformers import AutoConfig, AutoTokenizer
|
||||||
|
import signal
|
||||||
|
|
||||||
def check_vllm_status():
|
from .synthetic_configs import (
|
||||||
try:
|
synthetic_qa_config,
|
||||||
response = requests.get("http://localhost:8000/metrics")
|
)
|
||||||
if response.status_code == 200:
|
|
||||||
return True
|
class SyntheticDataKit:
|
||||||
except requests.exceptions.ConnectionError:
|
def __init__(
|
||||||
return False
|
self,
|
||||||
|
model_name = "unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
|
||||||
|
max_seq_length = 2048,
|
||||||
|
gpu_memory_utilization = 0.98,
|
||||||
|
float8_kv_cache = False,
|
||||||
|
conservativeness = 1.0,
|
||||||
|
token = None,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
assert(type(model_name) is str)
|
||||||
|
assert(type(max_seq_length) is int)
|
||||||
|
assert(type(gpu_memory_utilization) is float)
|
||||||
|
assert(type(float8_kv_cache) is bool)
|
||||||
|
assert(type(conservativeness) is float)
|
||||||
|
assert(token is None or type(token) is str)
|
||||||
|
|
||||||
|
self.model_name = model_name
|
||||||
|
self.max_seq_length = max_seq_length
|
||||||
|
|
||||||
|
self.config = AutoConfig.from_pretrained(
|
||||||
|
model_name,
|
||||||
|
token = token,
|
||||||
|
)
|
||||||
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
||||||
|
model_name,
|
||||||
|
token = token,
|
||||||
|
)
|
||||||
|
engine_args = load_vllm(
|
||||||
|
model_name = model_name,
|
||||||
|
config = self.config,
|
||||||
|
gpu_memory_utilization = gpu_memory_utilization,
|
||||||
|
max_seq_length = max_seq_length,
|
||||||
|
disable_log_stats = True,
|
||||||
|
float8_kv_cache = float8_kv_cache,
|
||||||
|
conservativeness = conservativeness,
|
||||||
|
return_args = True,
|
||||||
|
enable_lora = False,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
if "device" in engine_args: del engine_args["device"]
|
||||||
|
if "model" in engine_args: del engine_args["model"]
|
||||||
|
if "compilation_config" in engine_args: del engine_args["compilation_config"]
|
||||||
|
|
||||||
|
subprocess_commands = [
|
||||||
|
"vllm", "serve", str(model_name),
|
||||||
|
]
|
||||||
|
for key, value in engine_args.items():
|
||||||
|
flag = key.replace("_", "-")
|
||||||
|
which = str(value).lower().replace("torch.", "")
|
||||||
|
if which == "true":
|
||||||
|
# Ignore --enforce-eager True
|
||||||
|
subprocess_commands += ["--" + flag,]
|
||||||
|
elif which == "false":
|
||||||
|
# Ignore flag
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
subprocess_commands += ["--" + flag, which,]
|
||||||
|
pass
|
||||||
|
vllm_process = subprocess.Popen(
|
||||||
|
subprocess_commands,
|
||||||
|
stdout = subprocess.PIPE,
|
||||||
|
stderr = subprocess.PIPE,
|
||||||
|
start_new_session = True,
|
||||||
|
)
|
||||||
|
self.vllm_process = vllm_process
|
||||||
|
|
||||||
|
ready_message_part = b"Starting vLLM API server on"
|
||||||
|
ready = False
|
||||||
|
while vllm_process.poll() is None:
|
||||||
|
output = vllm_process.stdout.readline()
|
||||||
|
if not output:
|
||||||
|
print("Stdout stream ended before readiness message detected.")
|
||||||
|
break
|
||||||
|
output_str = output.decode('utf-8', errors='ignore').strip()
|
||||||
|
print(f"vLLM STDOUT: {output_str}")
|
||||||
|
if ready_message_part in output:
|
||||||
|
print(f"\n--- vLLM Server Ready (Detected: '{ready_message_part.decode()}') ---")
|
||||||
|
ready = True
|
||||||
|
break
|
||||||
|
pass
|
||||||
|
pass
|
||||||
|
if vllm_process is None:
|
||||||
|
raise RuntimeError("Unsloth: vllm_process failed to load!")
|
||||||
|
trial = 0
|
||||||
|
while not self.check_vllm_status():
|
||||||
|
if trial >= 100:
|
||||||
|
raise RuntimeError("Unsloth: vllm_process failed to load!")
|
||||||
|
trial += 1
|
||||||
|
time.sleep(1)
|
||||||
|
return
|
||||||
pass
|
pass
|
||||||
pass
|
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
def async_load_vllm(
|
def from_pretrained(
|
||||||
model_name = "unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
|
model_name = "unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
|
||||||
max_seq_length = 2048,
|
max_seq_length = 2048,
|
||||||
gpu_memory_utilization = 0.85,
|
gpu_memory_utilization = 0.9,
|
||||||
float8_kv_cache = False,
|
float8_kv_cache = False,
|
||||||
conservativeness = 1.0,
|
conservativeness = 1.0,
|
||||||
token = None,
|
token = None,
|
||||||
):
|
**kwargs,
|
||||||
config = AutoConfig.from_pretrained(
|
):
|
||||||
model_name,
|
return SyntheticDataKit(
|
||||||
token = token,
|
model_name = model_name,
|
||||||
)
|
max_seq_length = max_seq_length,
|
||||||
engine_args = load_vllm(
|
gpu_memory_utilization = gpu_memory_utilization,
|
||||||
model_name = model_name,
|
float8_kv_cache = float8_kv_cache,
|
||||||
config = config,
|
conservativeness = conservativeness,
|
||||||
gpu_memory_utilization = gpu_memory_utilization,
|
token = token,
|
||||||
max_seq_length = max_seq_length,
|
**kwargs,
|
||||||
disable_log_stats = True,
|
)
|
||||||
float8_kv_cache = float8_kv_cache,
|
|
||||||
conservativeness = conservativeness,
|
|
||||||
return_args = True,
|
|
||||||
enable_lora = False,
|
|
||||||
)
|
|
||||||
if "device" in engine_args: del engine_args["device"]
|
|
||||||
if "model" in engine_args: del engine_args["model"]
|
|
||||||
|
|
||||||
subprocess_commands = [
|
|
||||||
"vllm", "serve", str(model_name),
|
|
||||||
]
|
|
||||||
for key, value in engine_args.items():
|
|
||||||
flag = "--" + key.replace("_", "-")
|
|
||||||
which = str(value).lower().replace("torch.", "")
|
|
||||||
subprocess_commands += [flag, which,]
|
|
||||||
pass
|
pass
|
||||||
print(subprocess_commands)
|
|
||||||
vllm_process = subprocess.Popen(
|
@staticmethod
|
||||||
subprocess_commands,
|
def check_vllm_status():
|
||||||
stdout = subprocess.PIPE,
|
try:
|
||||||
stderr = subprocess.PIPE,
|
response = requests.get("http://localhost:8000/metrics")
|
||||||
start_new_session = True,
|
if response.status_code == 200:
|
||||||
)
|
return True
|
||||||
ready_message_part = b"Starting vLLM API server on"
|
except requests.exceptions.ConnectionError:
|
||||||
ready = False
|
return False
|
||||||
while vllm_process.poll() is None:
|
|
||||||
output = vllm_process.stdout.readline()
|
|
||||||
if not output:
|
|
||||||
print("Stdout stream ended before readiness message detected.")
|
|
||||||
break
|
|
||||||
output_str = output.decode('utf-8', errors='ignore').strip()
|
|
||||||
print(f"vLLM STDOUT: {output_str}")
|
|
||||||
if ready_message_part in output:
|
|
||||||
print(f"\n--- vLLM Server Ready (Detected: '{ready_message_part.decode()}') ---")
|
|
||||||
ready = True
|
|
||||||
break
|
|
||||||
pass
|
pass
|
||||||
pass
|
pass
|
||||||
if vllm_process is None:
|
|
||||||
raise RuntimeError("Unsloth: vllm_process failed to load!")
|
def cleanup(self):
|
||||||
trial = 0
|
if not hasattr(self, "vllm_process"): return
|
||||||
while not check_vllm_status():
|
|
||||||
if trial >= 100:
|
vllm_process = self.vllm_process
|
||||||
raise RuntimeError("Unsloth: vllm_process failed to load!")
|
print("Attempting to terminate the VLLM server gracefully...")
|
||||||
trial += 1
|
try:
|
||||||
time.sleep(1)
|
vllm_process.terminate()
|
||||||
return vllm_process
|
vllm_process.wait(timeout=10)
|
||||||
pass
|
print("Server terminated gracefully.")
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
print("Server did not terminate gracefully after 10 seconds. Forcing kill...")
|
||||||
def destroy_vllm(vllm_process):
|
vllm_process.kill()
|
||||||
print("Attempting to terminate the VLLM server gracefully...")
|
vllm_process.wait()
|
||||||
try:
|
print("Server killed forcefully.")
|
||||||
vllm_process.terminate()
|
except Exception as e:
|
||||||
vllm_process.wait(timeout=10)
|
print(f"An error occurred while trying to stop the process: {e}")
|
||||||
print("Server terminated gracefully.")
|
try:
|
||||||
except subprocess.TimeoutExpired:
|
if vllm_process.poll() is None:
|
||||||
print("Server did not terminate gracefully after 10 seconds. Forcing kill...")
|
print("Attempting forceful kill due to error...")
|
||||||
vllm_process.kill()
|
vllm_process.kill()
|
||||||
vllm_process.wait()
|
vllm_process.wait()
|
||||||
print("Server killed forcefully.")
|
print("Server killed forcefully after error.")
|
||||||
except Exception as e:
|
except Exception as kill_e:
|
||||||
print(f"An error occurred while trying to stop the process: {e}")
|
print(f"Error during forceful kill: {kill_e}")
|
||||||
try:
|
for _ in range(10):
|
||||||
if vllm_process.poll() is None:
|
torch.cuda.empty_cache()
|
||||||
print("Attempting forceful kill due to error...")
|
gc.collect()
|
||||||
vllm_process.kill()
|
pass
|
||||||
vllm_process.wait()
|
|
||||||
print("Server killed forcefully after error.")
|
def __enter__(self): return self
|
||||||
except Exception as kill_e:
|
def __exit__(self, *exc): self.cleanup()
|
||||||
print(f"Error during forceful kill: {kill_e}")
|
def __del__(self): self.cleanup()
|
||||||
for _ in range(10):
|
|
||||||
torch.cuda.empty_cache()
|
def truncate(self, filename = None):
|
||||||
gc.collect()
|
# Truncates by summary and max generation
|
||||||
pass
|
assert(filename is not None)
|
||||||
|
assert(os.path.exists(filename))
|
||||||
|
assert(hasattr(self, "tokenizer"))
|
||||||
synthetic_config_string = """\
|
|
||||||
# Master configuration file for Synthetic Data Kit
|
with open(filename, "r") as f: text = f.read()
|
||||||
|
|
||||||
# Global paths configuration
|
max_tokens = self.max_seq_length - self.max_generation_tokens*2 - 2
|
||||||
paths:
|
input_ids = self.tokenizer(text).input_ids
|
||||||
# Input data locations
|
length = len(text)
|
||||||
input:
|
original_length = len(text)
|
||||||
pdf: "data/pdf"
|
original_n_tokens = len(input_ids)
|
||||||
html: "data/html"
|
|
||||||
youtube: "data/youtube"
|
if len(input_ids) > max_tokens:
|
||||||
docx: "data/docx"
|
# Will fix later, but for now we simply naively truncate by ratios
|
||||||
ppt: "data/ppt"
|
length = original_length
|
||||||
txt: "data/txt"
|
while True:
|
||||||
|
input_ids = self.tokenizer(text[:length]).input_ids
|
||||||
# Output locations
|
if len(input_ids) < max_tokens or length == 0: break
|
||||||
output:
|
length = length * (max_tokens/len(input_ids))
|
||||||
parsed: "data/output" # Where parsed text files are saved
|
length = max(int(length), 0)
|
||||||
generated: "data/generated" # Where generated content is saved
|
pass
|
||||||
cleaned: "data/cleaned" # Where cleaned content is saved
|
print(f"Unsloth: Will truncate your data which has {original_n_tokens} tokens to {len(input_ids)} tokens.")
|
||||||
final: "data/final" # Where final formatted content is saved
|
|
||||||
|
with open(filename, "w") as f: f.write(text[:length])
|
||||||
# VLLM server configuration
|
pass
|
||||||
vllm:
|
return filename, length
|
||||||
api_base: "http://localhost:8000/v1" # Base URL for VLLM API
|
pass
|
||||||
port: 8000 # Port for VLLM server
|
|
||||||
model: "{model_name}" # Default model to use
|
def prepare_qa_generation(
|
||||||
max_retries: 3 # Number of retries for API calls
|
self,
|
||||||
retry_delay: 1.0 # Initial delay between retries (seconds)
|
output_folder = "data",
|
||||||
|
max_generation_tokens = 512,
|
||||||
# Ingest configuration
|
temperature = 0.7,
|
||||||
ingest:
|
top_p = 0.95,
|
||||||
default_format: "txt" # Default output format for parsed files
|
overlap = 64,
|
||||||
youtube_captions: "auto" # Options: "auto", "manual" - caption preference
|
default_num_pairs = 25,
|
||||||
|
cleanup_threshold = 1.0,
|
||||||
# LLM generation parameters
|
cleanup_batch_size = 4,
|
||||||
generation:
|
cleanup_temperature = 0.3,
|
||||||
temperature: {temperature} # Higher = more creative, lower = more deterministic
|
):
|
||||||
top_p: {top_p} # Nucleus sampling parameter
|
assert(hasattr(self, "model_name"))
|
||||||
chunk_size: {chunk_size} # Size of text chunks for processing
|
assert(hasattr(self, "max_seq_length"))
|
||||||
overlap: {overlap} # Overlap between chunks to maintain context
|
assert(max_generation_tokens < self.max_seq_length)
|
||||||
max_tokens: {max_tokens} # Maximum tokens in LLM responses
|
|
||||||
num_pairs: {default_num_pairs} # Default number of QA pairs to generate
|
locations = "pdf,html,youtube,docx,ppt,txt,output,generated,cleaned,final"
|
||||||
|
locations = locations.split(",")
|
||||||
# Content cleanup parameters
|
for path in locations:
|
||||||
cleanup:
|
os.makedirs(os.path.join(output_folder, path), exist_ok = True)
|
||||||
threshold: {cleanup_threshold} # Default quality threshold (1-10)
|
pass
|
||||||
batch_size: {cleanup_batch_size} # Number of items per batch for rating
|
|
||||||
temperature: {cleanup_temperature} # Temperature for rating (lower = more consistent)
|
self.max_generation_tokens = max_generation_tokens
|
||||||
|
|
||||||
# Format conversion parameters
|
config = synthetic_qa_config\
|
||||||
format:
|
.replace("{data_output_location}", str(output_folder))\
|
||||||
default: "jsonl" # Default output format
|
.replace("{model_name}", str(self.model_name))\
|
||||||
include_metadata: true # Include metadata in output files
|
.replace("{temperature}", str(temperature))\
|
||||||
pretty_json: true # Use indentation in JSON output
|
.replace("{top_p}", str(top_p))\
|
||||||
|
.replace("{chunk_size}", str(self.max_seq_length - max_generation_tokens*2 - 2))\
|
||||||
# Prompts for different tasks
|
.replace("{overlap}", str(overlap))\
|
||||||
prompts:
|
.replace("{max_tokens}", str(max_generation_tokens))\
|
||||||
# Summary generation prompt
|
.replace("{default_num_pairs}", str(default_num_pairs))\
|
||||||
summary: |
|
.replace("{cleanup_threshold}", str(cleanup_threshold))\
|
||||||
Summarize this document in 3-5 sentences, focusing on the main topic and key concepts.
|
.replace("{cleanup_batch_size}", str(cleanup_batch_size))\
|
||||||
|
.replace("{cleanup_temperature}", str(cleanup_temperature))
|
||||||
# QA pair generation prompt
|
|
||||||
qa_generation: |
|
with open("synthetic_data_kit_config.yaml", "w") as f: f.write(config)
|
||||||
Create {num_pairs} question-answer pairs from this text for LLM training.
|
pass
|
||||||
|
|
||||||
Rules:
|
|
||||||
1. Questions must be about important facts in the text
|
|
||||||
2. Answers must be directly supported by the text
|
|
||||||
3. Return JSON format only:
|
|
||||||
|
|
||||||
[
|
|
||||||
{{
|
|
||||||
"question": "Question 1?",
|
|
||||||
"answer": "Answer 1."
|
|
||||||
}},
|
|
||||||
{{
|
|
||||||
"question": "Question 2?",
|
|
||||||
"answer": "Answer 2."
|
|
||||||
}}
|
|
||||||
]
|
|
||||||
|
|
||||||
Text:
|
|
||||||
{text}
|
|
||||||
|
|
||||||
# QA pair rating prompt
|
|
||||||
qa_rating: |
|
|
||||||
Rate each of these question-answer pairs for quality and return exactly this JSON format:
|
|
||||||
|
|
||||||
[
|
|
||||||
{{"question": "same question text", "answer": "same answer text", "rating": n}}
|
|
||||||
]
|
|
||||||
|
|
||||||
Where n is a number from 1-10.
|
|
||||||
|
|
||||||
DO NOT include any text outside of the JSON array, just return valid JSON:
|
|
||||||
|
|
||||||
{pairs}"""
|
|
||||||
|
|
||||||
|
|
||||||
def configure_synthetic_data_kit(
|
|
||||||
model_name = "unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
|
|
||||||
temperature = 0.7,
|
|
||||||
top_p = 0.95,
|
|
||||||
chunk_size = 4000,
|
|
||||||
overlap = 200,
|
|
||||||
max_tokens = 512,
|
|
||||||
default_num_pairs = 25,
|
|
||||||
cleanup_threshold = 1.0,
|
|
||||||
cleanup_batch_size = 4,
|
|
||||||
cleanup_temperature = 0.3,
|
|
||||||
):
|
|
||||||
config = synthetic_config_string\
|
|
||||||
.replace("{model_name}", str(model_name))\
|
|
||||||
.replace("{temperature}", str(temperature))\
|
|
||||||
.replace("{top_p}", str(top_p))\
|
|
||||||
.replace("{chunk_size}", str(chunk_size))\
|
|
||||||
.replace("{overlap}", str(overlap))\
|
|
||||||
.replace("{max_tokens}", str(max_tokens))\
|
|
||||||
.replace("{default_num_pairs}", str(default_num_pairs))\
|
|
||||||
.replace("{cleanup_threshold}", str(cleanup_threshold))\
|
|
||||||
.replace("{cleanup_batch_size}", str(cleanup_batch_size))\
|
|
||||||
.replace("{cleanup_temperature}", str(cleanup_temperature))
|
|
||||||
|
|
||||||
return config
|
|
||||||
pass
|
pass
|
||||||
|
|
|
||||||
111
unsloth/dataprep/synthetic_configs.py
Normal file
111
unsloth/dataprep/synthetic_configs.py
Normal file
|
|
@ -0,0 +1,111 @@
|
||||||
|
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
synthetic_qa_config = """\
|
||||||
|
# Master configuration file for Synthetic Data Kit
|
||||||
|
|
||||||
|
# Global paths configuration
|
||||||
|
paths:
|
||||||
|
# Input data locations
|
||||||
|
input:
|
||||||
|
pdf: "{data_output_location}/pdf"
|
||||||
|
html: "{data_output_location}/html"
|
||||||
|
youtube: "{data_output_location}/youtube"
|
||||||
|
docx: "{data_output_location}/docx"
|
||||||
|
ppt: "{data_output_location}/ppt"
|
||||||
|
txt: "{data_output_location}/txt"
|
||||||
|
|
||||||
|
# Output locations
|
||||||
|
output:
|
||||||
|
parsed: "{data_output_location}/output" # Where parsed text files are saved
|
||||||
|
generated: "{data_output_location}/generated" # Where generated content is saved
|
||||||
|
cleaned: "{data_output_location}/cleaned" # Where cleaned content is saved
|
||||||
|
final: "{data_output_location}/final" # Where final formatted content is saved
|
||||||
|
|
||||||
|
# VLLM server configuration
|
||||||
|
vllm:
|
||||||
|
api_base: "http://localhost:8000/v1" # Base URL for VLLM API
|
||||||
|
port: 8000 # Port for VLLM server
|
||||||
|
model: "{model_name}" # Default model to use
|
||||||
|
max_retries: 3 # Number of retries for API calls
|
||||||
|
retry_delay: 1.0 # Initial delay between retries (seconds)
|
||||||
|
|
||||||
|
# Ingest configuration
|
||||||
|
ingest:
|
||||||
|
default_format: "txt" # Default output format for parsed files
|
||||||
|
youtube_captions: "auto" # Options: "auto", "manual" - caption preference
|
||||||
|
|
||||||
|
# LLM generation parameters
|
||||||
|
generation:
|
||||||
|
temperature: {temperature} # Higher = more creative, lower = more deterministic
|
||||||
|
top_p: {top_p} # Nucleus sampling parameter
|
||||||
|
chunk_size: {chunk_size} # Size of text chunks for processing
|
||||||
|
overlap: {overlap} # Overlap between chunks to maintain context
|
||||||
|
max_tokens: {max_tokens} # Maximum tokens in LLM responses
|
||||||
|
num_pairs: {default_num_pairs} # Default number of QA pairs to generate
|
||||||
|
|
||||||
|
# Content cleanup parameters
|
||||||
|
cleanup:
|
||||||
|
threshold: {cleanup_threshold} # Default quality threshold (1-10)
|
||||||
|
batch_size: {cleanup_batch_size} # Number of items per batch for rating
|
||||||
|
temperature: {cleanup_temperature} # Temperature for rating (lower = more consistent)
|
||||||
|
|
||||||
|
# Format conversion parameters
|
||||||
|
format:
|
||||||
|
default: "jsonl" # Default output format
|
||||||
|
include_metadata: true # Include metadata in output files
|
||||||
|
pretty_json: true # Use indentation in JSON output
|
||||||
|
|
||||||
|
# Prompts for different tasks
|
||||||
|
prompts:
|
||||||
|
# Summary generation prompt
|
||||||
|
summary: |
|
||||||
|
Summarize this document in 3-5 sentences, focusing on the main topic and key concepts.
|
||||||
|
|
||||||
|
# QA pair generation prompt
|
||||||
|
qa_generation: |
|
||||||
|
Create {num_pairs} question-answer pairs from this text for LLM training.
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
1. Questions must be about important facts in the text
|
||||||
|
2. Answers must be directly supported by the text
|
||||||
|
3. Return JSON format only:
|
||||||
|
|
||||||
|
[
|
||||||
|
{{
|
||||||
|
"question": "Question 1?",
|
||||||
|
"answer": "Answer 1."
|
||||||
|
}},
|
||||||
|
{{
|
||||||
|
"question": "Question 2?",
|
||||||
|
"answer": "Answer 2."
|
||||||
|
}}
|
||||||
|
]
|
||||||
|
|
||||||
|
Text:
|
||||||
|
{text}
|
||||||
|
|
||||||
|
# QA pair rating prompt
|
||||||
|
qa_rating: |
|
||||||
|
Rate each of these question-answer pairs for quality and return exactly this JSON format:
|
||||||
|
|
||||||
|
[
|
||||||
|
{{"question": "same question text", "answer": "same answer text", "rating": n}}
|
||||||
|
]
|
||||||
|
|
||||||
|
Where n is a number from 1-10.
|
||||||
|
|
||||||
|
DO NOT include any text outside of the JSON array, just return valid JSON:
|
||||||
|
|
||||||
|
{pairs}"""
|
||||||
16
unsloth/models/llama4.py
Normal file
16
unsloth/models/llama4.py
Normal file
|
|
@ -0,0 +1,16 @@
|
||||||
|
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from unsloth_studio.models import patch_llama4
|
||||||
|
patch_llama4()
|
||||||
|
|
@ -302,7 +302,7 @@ class FastLanguageModel(FastLlamaModel):
|
||||||
dispatch_model = FastGemma2Model
|
dispatch_model = FastGemma2Model
|
||||||
elif model_type == "qwen2":
|
elif model_type == "qwen2":
|
||||||
dispatch_model = FastQwen2Model
|
dispatch_model = FastQwen2Model
|
||||||
elif model_type == "qwen3" or model_type == "qwen3_moe":
|
elif model_type == "qwen3":# or model_type == "qwen3_moe":
|
||||||
if not SUPPORTS_QWEN3 or not SUPPORTS_QWEN3_MOE:
|
if not SUPPORTS_QWEN3 or not SUPPORTS_QWEN3_MOE:
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
f"Unsloth: Your transformers version of {transformers_version} does not support Qwen3.\n"\
|
f"Unsloth: Your transformers version of {transformers_version} does not support Qwen3.\n"\
|
||||||
|
|
|
||||||
|
|
@ -797,15 +797,6 @@ __INT_TO_FLOAT_MAPPER = \
|
||||||
"Qwen/Qwen3-14B-Base",
|
"Qwen/Qwen3-14B-Base",
|
||||||
"unsloth/Qwen3-14B-Base-bnb-4bit",
|
"unsloth/Qwen3-14B-Base-bnb-4bit",
|
||||||
),
|
),
|
||||||
"unsloth/Qwen3-32B-Base-unsloth-bnb-4bit" : (
|
|
||||||
"unsloth/Qwen3-32B-Base",
|
|
||||||
"Qwen/Qwen3-32B-Base",
|
|
||||||
"unsloth/Qwen3-32B-Base-bnb-4bit",
|
|
||||||
),
|
|
||||||
"unsloth/Qwen3-30B-A3B-Base-bnb-4bit" : (
|
|
||||||
"unsloth/Qwen3-30B-A3B-Base",
|
|
||||||
"Qwen/Qwen3-30B-A3B-Base",
|
|
||||||
),
|
|
||||||
}
|
}
|
||||||
|
|
||||||
INT_TO_FLOAT_MAPPER = {}
|
INT_TO_FLOAT_MAPPER = {}
|
||||||
|
|
|
||||||
110
unsloth/registry/REGISTRY.md
Normal file
110
unsloth/registry/REGISTRY.md
Normal file
|
|
@ -0,0 +1,110 @@
|
||||||
|
## Model Registry
|
||||||
|
|
||||||
|
### Structure
|
||||||
|
```
|
||||||
|
unsloth
|
||||||
|
-registry
|
||||||
|
__init__.py
|
||||||
|
registry.py
|
||||||
|
_llama.py
|
||||||
|
_mistral.py
|
||||||
|
_phi.py
|
||||||
|
...
|
||||||
|
```
|
||||||
|
|
||||||
|
Each model is registered in a separate file within the `registry` module (e.g. `registry/_llama.py`).
|
||||||
|
|
||||||
|
Within each model registration file, a high-level `ModelMeta` is created for each model version, with the following structure:
|
||||||
|
```python
|
||||||
|
@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[QuantType] | dict[str, list[QuantType]] = field(default_factory=list)
|
||||||
|
is_multimodal: bool = False
|
||||||
|
```
|
||||||
|
|
||||||
|
Each model then instantiates a global `ModelMeta` for its specific model version, defining how the model path (e.g. `unsloth/Llama-3.1-8B-Instruct`) is constructed since each model type has a different naming convention.
|
||||||
|
```python
|
||||||
|
LlamaMeta_3_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=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
`LlamaModelInfo` is a subclass of `ModelInfo` that defines the model path for each model size and quant type.
|
||||||
|
```python
|
||||||
|
class LlamaModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}-{version}-{size}B"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
```
|
||||||
|
|
||||||
|
Once these constructs are defined, the model is registered by writing a register_xx_models function.
|
||||||
|
```python
|
||||||
|
def register_llama_3_1_models(include_original_model: bool = False):
|
||||||
|
global _IS_LLAMA_3_1_REGISTERED
|
||||||
|
if _IS_LLAMA_3_1_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(LlamaMeta_3_1, include_original_model=include_original_model)
|
||||||
|
_IS_LLAMA_3_1_REGISTERED = True
|
||||||
|
```
|
||||||
|
|
||||||
|
`_register_models` is a helper function that registers the model with the registry. The global `_IS_XX_REGISTERED` is used to prevent duplicate registration.
|
||||||
|
|
||||||
|
Once a model is registered, registry.registry.MODEL_REGISTRY is updated with the model info and can be searched with `registry.search_models`.
|
||||||
|
|
||||||
|
### Tests
|
||||||
|
|
||||||
|
The `tests/test_model_registry.py` file contains tests for the model registry.
|
||||||
|
|
||||||
|
Also, each model registration file is an executable module that checks that all registered models are available on `huggingface_hub`.
|
||||||
|
```python
|
||||||
|
python unsloth.registry._llama.py
|
||||||
|
```
|
||||||
|
|
||||||
|
Prints the following (abridged) output:
|
||||||
|
```bash
|
||||||
|
✓ unsloth/Llama-3.1-8B
|
||||||
|
✓ unsloth/Llama-3.1-8B-bnb-4bit
|
||||||
|
✓ unsloth/Llama-3.1-8B-unsloth-bnb-4bit
|
||||||
|
✓ meta-llama/Llama-3.1-8B
|
||||||
|
✓ unsloth/Llama-3.1-8B-Instruct
|
||||||
|
✓ unsloth/Llama-3.1-8B-Instruct-bnb-4bit
|
||||||
|
✓ unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit
|
||||||
|
✓ meta-llama/Llama-3.1-8B-Instruct
|
||||||
|
✓ unsloth/Llama-3.2-1B
|
||||||
|
✓ unsloth/Llama-3.2-1B-bnb-4bit
|
||||||
|
✓ unsloth/Llama-3.2-1B-unsloth-bnb-4bit
|
||||||
|
✓ meta-llama/Llama-3.2-1B
|
||||||
|
...
|
||||||
|
```
|
||||||
|
|
||||||
|
### TODO
|
||||||
|
- Model Collections
|
||||||
|
- [x] Gemma3
|
||||||
|
- [ ] Llama3.1
|
||||||
|
- [x] Llama3.2
|
||||||
|
- [x] MistralSmall
|
||||||
|
- [x] Qwen2.5
|
||||||
|
- [x] Qwen2.5-VL
|
||||||
|
- [ ] Qwen2.5 Coder
|
||||||
|
- [x] QwenQwQ-32B
|
||||||
|
- [x] Deepseek v3
|
||||||
|
- [x] Deepseek R1
|
||||||
|
- [x] Phi-4
|
||||||
|
- [ ] Unsloth 4-bit Dynamic Quants
|
||||||
|
- [ ] Vision/multimodal models
|
||||||
|
- Sync model uploads with registry
|
||||||
|
- Add utility methods for tracking model stats
|
||||||
51
unsloth/registry/__init__.py
Normal file
51
unsloth/registry/__init__.py
Normal file
|
|
@ -0,0 +1,51 @@
|
||||||
|
from ._deepseek import register_deepseek_models as _register_deepseek_models
|
||||||
|
from ._gemma import register_gemma_models as _register_gemma_models
|
||||||
|
from ._llama import register_llama_models as _register_llama_models
|
||||||
|
from ._mistral import register_mistral_models as _register_mistral_models
|
||||||
|
from ._phi import register_phi_models as _register_phi_models
|
||||||
|
from ._qwen import register_qwen_models as _register_qwen_models
|
||||||
|
from .registry import MODEL_REGISTRY, ModelInfo, QuantType
|
||||||
|
|
||||||
|
_ARE_MODELS_REGISTERED = False
|
||||||
|
|
||||||
|
def register_models():
|
||||||
|
global _ARE_MODELS_REGISTERED
|
||||||
|
|
||||||
|
if _ARE_MODELS_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_deepseek_models()
|
||||||
|
_register_gemma_models()
|
||||||
|
_register_llama_models()
|
||||||
|
_register_mistral_models()
|
||||||
|
_register_phi_models()
|
||||||
|
_register_qwen_models()
|
||||||
|
|
||||||
|
_ARE_MODELS_REGISTERED = True
|
||||||
|
|
||||||
|
def search_models(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]:
|
||||||
|
"""
|
||||||
|
Get model info from the registry.
|
||||||
|
|
||||||
|
See registry.ModelInfo for more fields.
|
||||||
|
|
||||||
|
If search_pattern is provided, the full model path will be matched against the pattern, where the model path is the model_id on huggingface hub.
|
||||||
|
|
||||||
|
"""
|
||||||
|
if not _ARE_MODELS_REGISTERED:
|
||||||
|
register_models()
|
||||||
|
|
||||||
|
model_infos = MODEL_REGISTRY.values()
|
||||||
|
if org:
|
||||||
|
model_infos = [model_info for model_info in model_infos if model_info.org == org]
|
||||||
|
if base_name:
|
||||||
|
model_infos = [model_info for model_info in model_infos if model_info.base_name == base_name]
|
||||||
|
if version:
|
||||||
|
model_infos = [model_info for model_info in model_infos if model_info.version == version]
|
||||||
|
if size:
|
||||||
|
model_infos = [model_info for model_info in model_infos if model_info.size == size]
|
||||||
|
if quant_types:
|
||||||
|
model_infos = [model_info for model_info in model_infos if any(model_info.quant_type == quant_type for quant_type in quant_types)]
|
||||||
|
if search_pattern:
|
||||||
|
model_infos = [model_info for model_info in model_infos if search_pattern in model_info.model_path]
|
||||||
|
|
||||||
|
return model_infos
|
||||||
179
unsloth/registry/_deepseek.py
Normal file
179
unsloth/registry/_deepseek.py
Normal file
|
|
@ -0,0 +1,179 @@
|
||||||
|
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
|
||||||
|
|
||||||
|
_IS_DEEPSEEK_V3_REGISTERED = False
|
||||||
|
_IS_DEEPSEEK_V3_0324_REGISTERED = False
|
||||||
|
_IS_DEEPSEEK_R1_REGISTERED = False
|
||||||
|
_IS_DEEPSEEK_R1_ZERO_REGISTERED = False
|
||||||
|
_IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED = False
|
||||||
|
_IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED = False
|
||||||
|
|
||||||
|
class DeepseekV3ModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}-V{version}"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, 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
|
||||||
|
if size:
|
||||||
|
key = f"{key}-{size}B"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
# Deepseek V3 Model Meta
|
||||||
|
DeepseekV3Meta = ModelMeta(
|
||||||
|
org="deepseek-ai",
|
||||||
|
base_name="DeepSeek",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="3",
|
||||||
|
model_sizes=[""],
|
||||||
|
model_info_cls=DeepseekV3ModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BF16],
|
||||||
|
)
|
||||||
|
|
||||||
|
DeepseekV3_0324Meta = ModelMeta(
|
||||||
|
org="deepseek-ai",
|
||||||
|
base_name="DeepSeek",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="3-0324",
|
||||||
|
model_sizes=[""],
|
||||||
|
model_info_cls=DeepseekV3ModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
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],
|
||||||
|
)
|
||||||
|
|
||||||
|
DeepseekR1ZeroMeta = ModelMeta(
|
||||||
|
org="deepseek-ai",
|
||||||
|
base_name="DeepSeek-R1",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="Zero",
|
||||||
|
model_sizes=[""],
|
||||||
|
model_info_cls=DeepseekR1ModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.GGUF],
|
||||||
|
)
|
||||||
|
|
||||||
|
DeepseekR1DistillLlamaMeta = 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]},
|
||||||
|
)
|
||||||
|
|
||||||
|
# Deepseek R1 Distill Qwen Model Meta
|
||||||
|
DeepseekR1DistillQwenMeta = ModelMeta(
|
||||||
|
org="deepseek-ai",
|
||||||
|
base_name="DeepSeek-R1-Distill",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="Qwen",
|
||||||
|
model_sizes=["1.5", "7", "14", "32"],
|
||||||
|
model_info_cls=DeepseekR1ModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types={
|
||||||
|
"1.5": [QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF],
|
||||||
|
"7": [QuantType.UNSLOTH, QuantType.BNB],
|
||||||
|
"14": [QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF],
|
||||||
|
"32": [QuantType.GGUF, QuantType.BNB],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
def register_deepseek_v3_models(include_original_model: bool = False):
|
||||||
|
global _IS_DEEPSEEK_V3_REGISTERED
|
||||||
|
if _IS_DEEPSEEK_V3_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(DeepseekV3Meta, include_original_model=include_original_model)
|
||||||
|
_IS_DEEPSEEK_V3_REGISTERED = True
|
||||||
|
|
||||||
|
def register_deepseek_v3_0324_models(include_original_model: bool = False):
|
||||||
|
global _IS_DEEPSEEK_V3_0324_REGISTERED
|
||||||
|
if _IS_DEEPSEEK_V3_0324_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(DeepseekV3_0324Meta, include_original_model=include_original_model)
|
||||||
|
_IS_DEEPSEEK_V3_0324_REGISTERED = True
|
||||||
|
|
||||||
|
def register_deepseek_r1_models(include_original_model: bool = False):
|
||||||
|
global _IS_DEEPSEEK_R1_REGISTERED
|
||||||
|
if _IS_DEEPSEEK_R1_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(DeepseekR1Meta, include_original_model=include_original_model)
|
||||||
|
_IS_DEEPSEEK_R1_REGISTERED = True
|
||||||
|
|
||||||
|
def register_deepseek_r1_zero_models(include_original_model: bool = False):
|
||||||
|
global _IS_DEEPSEEK_R1_ZERO_REGISTERED
|
||||||
|
if _IS_DEEPSEEK_R1_ZERO_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model)
|
||||||
|
_IS_DEEPSEEK_R1_ZERO_REGISTERED = True
|
||||||
|
|
||||||
|
def register_deepseek_r1_distill_llama_models(include_original_model: bool = False):
|
||||||
|
global _IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED
|
||||||
|
if _IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(DeepseekR1DistillLlamaMeta, include_original_model=include_original_model)
|
||||||
|
_IS_DEEPSEEK_R1_DISTILL_LLAMA_REGISTERED = True
|
||||||
|
|
||||||
|
def register_deepseek_r1_distill_qwen_models(include_original_model: bool = False):
|
||||||
|
global _IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED
|
||||||
|
if _IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(DeepseekR1DistillQwenMeta, include_original_model=include_original_model)
|
||||||
|
_IS_DEEPSEEK_R1_DISTILL_QWEN_REGISTERED = True
|
||||||
|
|
||||||
|
def register_deepseek_models(include_original_model: bool = False):
|
||||||
|
register_deepseek_v3_models(include_original_model=include_original_model)
|
||||||
|
register_deepseek_v3_0324_models(include_original_model=include_original_model)
|
||||||
|
register_deepseek_r1_models(include_original_model=include_original_model)
|
||||||
|
register_deepseek_r1_zero_models(include_original_model=include_original_model)
|
||||||
|
register_deepseek_r1_distill_llama_models(include_original_model=include_original_model)
|
||||||
|
register_deepseek_r1_distill_qwen_models(include_original_model=include_original_model)
|
||||||
|
|
||||||
|
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", limit=1000)
|
||||||
|
distill_models = []
|
||||||
|
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-")
|
||||||
|
distill_models.append(version)
|
||||||
|
|
||||||
|
return distill_models
|
||||||
|
|
||||||
|
|
||||||
|
register_deepseek_models(include_original_model=True)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
|
||||||
|
register_deepseek_models(include_original_model=True)
|
||||||
|
|
||||||
|
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}")
|
||||||
|
# distill_models = _list_deepseek_r1_distill_models()
|
||||||
|
# for model in sorted(distill_models):
|
||||||
|
# if "qwen" in model.lower():
|
||||||
|
# print(model)
|
||||||
66
unsloth/registry/_gemma.py
Normal file
66
unsloth/registry/_gemma.py
Normal file
|
|
@ -0,0 +1,66 @@
|
||||||
|
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
|
||||||
|
|
||||||
|
_IS_GEMMA_3_BASE_REGISTERED = False
|
||||||
|
_IS_GEMMA_3_INSTRUCT_REGISTERED = False
|
||||||
|
|
||||||
|
class GemmaModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}-{version}-{size}B"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
# Gemma3 Base Model Meta
|
||||||
|
GemmaMeta3Base = ModelMeta(
|
||||||
|
org="google",
|
||||||
|
base_name="gemma",
|
||||||
|
instruct_tags=["pt"], # pt = base
|
||||||
|
model_version="3",
|
||||||
|
model_sizes=["1", "4", "12", "27"],
|
||||||
|
model_info_cls=GemmaModelInfo,
|
||||||
|
is_multimodal=True,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Gemma3 Instruct Model Meta
|
||||||
|
GemmaMeta3Instruct = ModelMeta(
|
||||||
|
org="google",
|
||||||
|
base_name="gemma",
|
||||||
|
instruct_tags=["it"], # it = instruction tuned
|
||||||
|
model_version="3",
|
||||||
|
model_sizes=["1", "4", "12", "27"],
|
||||||
|
model_info_cls=GemmaModelInfo,
|
||||||
|
is_multimodal=True,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
|
||||||
|
)
|
||||||
|
|
||||||
|
def register_gemma_3_base_models(include_original_model: bool = False):
|
||||||
|
global _IS_GEMMA_3_BASE_REGISTERED
|
||||||
|
if _IS_GEMMA_3_BASE_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(GemmaMeta3Base, include_original_model=include_original_model)
|
||||||
|
_IS_GEMMA_3_BASE_REGISTERED = True
|
||||||
|
|
||||||
|
def register_gemma_3_instruct_models(include_original_model: bool = False):
|
||||||
|
global _IS_GEMMA_3_INSTRUCT_REGISTERED
|
||||||
|
if _IS_GEMMA_3_INSTRUCT_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(GemmaMeta3Instruct, include_original_model=include_original_model)
|
||||||
|
_IS_GEMMA_3_INSTRUCT_REGISTERED = True
|
||||||
|
|
||||||
|
def register_gemma_models(include_original_model: bool = False):
|
||||||
|
register_gemma_3_base_models(include_original_model=include_original_model)
|
||||||
|
register_gemma_3_instruct_models(include_original_model=include_original_model)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
|
||||||
|
register_gemma_models(include_original_model=True)
|
||||||
|
|
||||||
|
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}")
|
||||||
113
unsloth/registry/_llama.py
Normal file
113
unsloth/registry/_llama.py
Normal file
|
|
@ -0,0 +1,113 @@
|
||||||
|
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
|
||||||
|
|
||||||
|
_IS_LLAMA_3_1_REGISTERED = False
|
||||||
|
_IS_LLAMA_3_2_REGISTERED = False
|
||||||
|
_IS_LLAMA_3_2_VISION_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"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, 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"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
|
||||||
|
# Llama 3.1
|
||||||
|
LlamaMeta_3_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=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Llama 3.2 Base Models
|
||||||
|
LlamaMeta_3_2_Base = ModelMeta(
|
||||||
|
org="meta-llama",
|
||||||
|
base_name="Llama",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="3.2",
|
||||||
|
model_sizes=["1", "3"],
|
||||||
|
model_info_cls=LlamaModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Llama 3.2 Instruction Tuned Models
|
||||||
|
LlamaMeta_3_2_Instruct = ModelMeta(
|
||||||
|
org="meta-llama",
|
||||||
|
base_name="Llama",
|
||||||
|
instruct_tags=["Instruct"],
|
||||||
|
model_version="3.2",
|
||||||
|
model_sizes=["1", "3"],
|
||||||
|
model_info_cls=LlamaModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Llama 3.2 Vision
|
||||||
|
LlamaMeta_3_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={
|
||||||
|
"11": [QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
"90": [QuantType.NONE],
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def register_llama_3_1_models(include_original_model: bool = False):
|
||||||
|
global _IS_LLAMA_3_1_REGISTERED
|
||||||
|
if _IS_LLAMA_3_1_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(LlamaMeta_3_1, include_original_model=include_original_model)
|
||||||
|
_IS_LLAMA_3_1_REGISTERED = True
|
||||||
|
|
||||||
|
def register_llama_3_2_models(include_original_model: bool = False):
|
||||||
|
global _IS_LLAMA_3_2_REGISTERED
|
||||||
|
if _IS_LLAMA_3_2_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(LlamaMeta_3_2_Base, include_original_model=include_original_model)
|
||||||
|
_register_models(LlamaMeta_3_2_Instruct, include_original_model=include_original_model)
|
||||||
|
_IS_LLAMA_3_2_REGISTERED = True
|
||||||
|
|
||||||
|
def register_llama_3_2_vision_models(include_original_model: bool = False):
|
||||||
|
global _IS_LLAMA_3_2_VISION_REGISTERED
|
||||||
|
if _IS_LLAMA_3_2_VISION_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(LlamaMeta_3_2_Vision, include_original_model=include_original_model)
|
||||||
|
_IS_LLAMA_3_2_VISION_REGISTERED = True
|
||||||
|
|
||||||
|
|
||||||
|
def register_llama_models(include_original_model: bool = False):
|
||||||
|
register_llama_3_1_models(include_original_model=include_original_model)
|
||||||
|
register_llama_3_2_models(include_original_model=include_original_model)
|
||||||
|
register_llama_3_2_vision_models(include_original_model=include_original_model)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
|
||||||
|
register_llama_models(include_original_model=True)
|
||||||
|
|
||||||
|
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}")
|
||||||
70
unsloth/registry/_mistral.py
Normal file
70
unsloth/registry/_mistral.py
Normal file
|
|
@ -0,0 +1,70 @@
|
||||||
|
import copy
|
||||||
|
|
||||||
|
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
|
||||||
|
|
||||||
|
_IS_MISTRAL_SMALL_REGISTERED = False
|
||||||
|
|
||||||
|
_MISTRAL_SMALL_03_25_VERSION = "2503"
|
||||||
|
_MISTRAL_SMALL_01_25_VERSION = "2501"
|
||||||
|
_MISTRAL_SMALL_09_24_VERSION = "2409" # Not uploaded to unsloth
|
||||||
|
|
||||||
|
class MistralSmallModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
if version == _MISTRAL_SMALL_03_25_VERSION:
|
||||||
|
key = f"{base_name}-3.1-{size}B-{instruct_tag}"
|
||||||
|
else:
|
||||||
|
key = f"{base_name}-{size}B-{instruct_tag}"
|
||||||
|
key += f"-{version}"
|
||||||
|
key = cls.append_quant_type(key, quant_type)
|
||||||
|
|
||||||
|
return key
|
||||||
|
|
||||||
|
|
||||||
|
MistralSmall_2503_Base_Meta = ModelMeta(
|
||||||
|
org="mistralai",
|
||||||
|
base_name="Mistral-Small",
|
||||||
|
instruct_tags=["Base"],
|
||||||
|
model_version=_MISTRAL_SMALL_03_25_VERSION,
|
||||||
|
model_sizes=["24"],
|
||||||
|
model_info_cls=MistralSmallModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB],
|
||||||
|
)
|
||||||
|
|
||||||
|
MistralSmall_2503_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Base_Meta)
|
||||||
|
MistralSmall_2503_Instruct_Meta.instruct_tags = ["Instruct"]
|
||||||
|
MistralSmall_2503_Instruct_Meta.quant_types = [QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF]
|
||||||
|
|
||||||
|
MistralSmall_2501_Base_Meta = copy.deepcopy(MistralSmall_2503_Base_Meta)
|
||||||
|
MistralSmall_2501_Base_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION
|
||||||
|
|
||||||
|
MistralSmall_2501_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Instruct_Meta)
|
||||||
|
MistralSmall_2501_Instruct_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION
|
||||||
|
|
||||||
|
def register_mistral_small_models(include_original_model: bool = False):
|
||||||
|
global _IS_MISTRAL_SMALL_REGISTERED
|
||||||
|
if _IS_MISTRAL_SMALL_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(MistralSmall_2503_Base_Meta, include_original_model=include_original_model)
|
||||||
|
_register_models(MistralSmall_2503_Instruct_Meta, include_original_model=include_original_model)
|
||||||
|
_register_models(MistralSmall_2501_Base_Meta, include_original_model=include_original_model)
|
||||||
|
_register_models(MistralSmall_2501_Instruct_Meta, include_original_model=include_original_model)
|
||||||
|
|
||||||
|
_IS_MISTRAL_SMALL_REGISTERED = True
|
||||||
|
|
||||||
|
def register_mistral_models(include_original_model: bool = False):
|
||||||
|
register_mistral_small_models(include_original_model=include_original_model)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
|
||||||
|
register_mistral_models(include_original_model=True)
|
||||||
|
|
||||||
|
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}")
|
||||||
65
unsloth/registry/_phi.py
Normal file
65
unsloth/registry/_phi.py
Normal file
|
|
@ -0,0 +1,65 @@
|
||||||
|
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
|
||||||
|
|
||||||
|
_IS_PHI_4_REGISTERED = False
|
||||||
|
_IS_PHI_4_INSTRUCT_REGISTERED = False
|
||||||
|
|
||||||
|
class PhiModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}-{version}"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
# Phi Model Meta
|
||||||
|
PhiMeta4 = ModelMeta(
|
||||||
|
org="microsoft",
|
||||||
|
base_name="phi",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="4",
|
||||||
|
model_sizes=["1"], # Assuming only one size
|
||||||
|
model_info_cls=PhiModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Phi Instruct Model Meta
|
||||||
|
PhiInstructMeta4 = ModelMeta(
|
||||||
|
org="microsoft",
|
||||||
|
base_name="phi",
|
||||||
|
instruct_tags=["mini-instruct"],
|
||||||
|
model_version="4",
|
||||||
|
model_sizes=["1"], # Assuming only one size
|
||||||
|
model_info_cls=PhiModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
|
||||||
|
)
|
||||||
|
|
||||||
|
def register_phi_4_models(include_original_model: bool = False):
|
||||||
|
global _IS_PHI_4_REGISTERED
|
||||||
|
if _IS_PHI_4_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(PhiMeta4, include_original_model=include_original_model)
|
||||||
|
_IS_PHI_4_REGISTERED = True
|
||||||
|
|
||||||
|
def register_phi_4_instruct_models(include_original_model: bool = False):
|
||||||
|
global _IS_PHI_4_INSTRUCT_REGISTERED
|
||||||
|
if _IS_PHI_4_INSTRUCT_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(PhiInstructMeta4, include_original_model=include_original_model)
|
||||||
|
_IS_PHI_4_INSTRUCT_REGISTERED = True
|
||||||
|
|
||||||
|
def register_phi_models(include_original_model: bool = False):
|
||||||
|
register_phi_4_models(include_original_model=include_original_model)
|
||||||
|
register_phi_4_instruct_models(include_original_model=include_original_model)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
|
||||||
|
register_phi_models(include_original_model=True)
|
||||||
|
|
||||||
|
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}")
|
||||||
117
unsloth/registry/_qwen.py
Normal file
117
unsloth/registry/_qwen.py
Normal file
|
|
@ -0,0 +1,117 @@
|
||||||
|
from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
|
||||||
|
|
||||||
|
_IS_QWEN_2_5_REGISTERED = False
|
||||||
|
_IS_QWEN_2_5_VL_REGISTERED = False
|
||||||
|
_IS_QWEN_QWQ_REGISTERED = False
|
||||||
|
class QwenModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}{version}-{size}B"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, 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"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
class QwenQwQModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}-{size}B"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
class QwenQVQPreviewModelInfo(ModelInfo):
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||||
|
key = f"{base_name}-{size}B-Preview"
|
||||||
|
return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
|
||||||
|
|
||||||
|
# Qwen2.5 Model Meta
|
||||||
|
Qwen_2_5_Meta = ModelMeta(
|
||||||
|
org="Qwen",
|
||||||
|
base_name="Qwen",
|
||||||
|
instruct_tags=[None, "Instruct"],
|
||||||
|
model_version="2.5",
|
||||||
|
model_sizes=["3", "7"],
|
||||||
|
model_info_cls=QwenModelInfo,
|
||||||
|
is_multimodal=False,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Qwen2.5 VL Model Meta
|
||||||
|
Qwen_2_5_VLMeta = ModelMeta(
|
||||||
|
org="Qwen",
|
||||||
|
base_name="Qwen",
|
||||||
|
instruct_tags=["Instruct"], # No base, only instruction tuned
|
||||||
|
model_version="2.5",
|
||||||
|
model_sizes=["3", "7", "32", "72"],
|
||||||
|
model_info_cls=QwenVLModelInfo,
|
||||||
|
is_multimodal=True,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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],
|
||||||
|
)
|
||||||
|
|
||||||
|
# Qwen QVQ Preview Model Meta
|
||||||
|
QwenQVQPreviewMeta = ModelMeta(
|
||||||
|
org="Qwen",
|
||||||
|
base_name="QVQ",
|
||||||
|
instruct_tags=[None],
|
||||||
|
model_version="",
|
||||||
|
model_sizes=["72"],
|
||||||
|
model_info_cls=QwenQVQPreviewModelInfo,
|
||||||
|
is_multimodal=True,
|
||||||
|
quant_types=[QuantType.NONE, QuantType.BNB],
|
||||||
|
)
|
||||||
|
|
||||||
|
def register_qwen_2_5_models(include_original_model: bool = False):
|
||||||
|
global _IS_QWEN_2_5_REGISTERED
|
||||||
|
if _IS_QWEN_2_5_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(Qwen_2_5_Meta, include_original_model=include_original_model)
|
||||||
|
_IS_QWEN_2_5_REGISTERED = True
|
||||||
|
|
||||||
|
def register_qwen_2_5_vl_models(include_original_model: bool = False):
|
||||||
|
global _IS_QWEN_2_5_VL_REGISTERED
|
||||||
|
if _IS_QWEN_2_5_VL_REGISTERED:
|
||||||
|
return
|
||||||
|
_register_models(Qwen_2_5_VLMeta, include_original_model=include_original_model)
|
||||||
|
_IS_QWEN_2_5_VL_REGISTERED = True
|
||||||
|
|
||||||
|
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=include_original_model)
|
||||||
|
_register_models(QwenQVQPreviewMeta, include_original_model=include_original_model)
|
||||||
|
_IS_QWEN_QWQ_REGISTERED = True
|
||||||
|
|
||||||
|
def register_qwen_models(include_original_model: bool = False):
|
||||||
|
register_qwen_2_5_models(include_original_model=include_original_model)
|
||||||
|
register_qwen_2_5_vl_models(include_original_model=include_original_model)
|
||||||
|
register_qwen_qwq_models(include_original_model=include_original_model)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
|
||||||
|
MODEL_REGISTRY.clear()
|
||||||
|
|
||||||
|
register_qwen_models(include_original_model=True)
|
||||||
|
|
||||||
|
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}")
|
||||||
185
unsloth/registry/registry.py
Normal file
185
unsloth/registry/registry.py
Normal file
|
|
@ -0,0 +1,185 @@
|
||||||
|
import warnings
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from enum import Enum
|
||||||
|
|
||||||
|
|
||||||
|
class QuantType(Enum):
|
||||||
|
BNB = "bnb"
|
||||||
|
UNSLOTH = "unsloth" # dynamic 4-bit quantization
|
||||||
|
GGUF = "GGUF"
|
||||||
|
NONE = "none"
|
||||||
|
BF16 = "bf16" # only for Deepseek V3
|
||||||
|
|
||||||
|
# Tags for Hugging Face model paths
|
||||||
|
BNB_QUANTIZED_TAG = "bnb-4bit"
|
||||||
|
UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG
|
||||||
|
GGUF_TAG = "GGUF"
|
||||||
|
BF16_TAG = "bf16"
|
||||||
|
|
||||||
|
QUANT_TAG_MAP = {
|
||||||
|
QuantType.BNB: BNB_QUANTIZED_TAG,
|
||||||
|
QuantType.UNSLOTH: UNSLOTH_DYNAMIC_QUANT_TAG,
|
||||||
|
QuantType.GGUF: GGUF_TAG,
|
||||||
|
QuantType.NONE: None,
|
||||||
|
QuantType.BF16: BF16_TAG,
|
||||||
|
}
|
||||||
|
|
||||||
|
# NOTE: models registered with org="unsloth" and QUANT_TYPE.NONE are aliases of QUANT_TYPE.UNSLOTH
|
||||||
|
@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: QuantType = None
|
||||||
|
description: str = 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: QuantType = None
|
||||||
|
):
|
||||||
|
if quant_type != QuantType.NONE:
|
||||||
|
key = "-".join([key, QUANT_TAG_MAP[quant_type]])
|
||||||
|
return key
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag, key=""):
|
||||||
|
key = cls.append_instruct_tag(key, instruct_tag)
|
||||||
|
key = cls.append_quant_type(key, quant_type)
|
||||||
|
return key
|
||||||
|
|
||||||
|
@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[QuantType] | dict[str, list[QuantType]] = 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: QuantType = 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, current keys: {MODEL_REGISTRY.keys()}")
|
||||||
|
|
||||||
|
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
|
||||||
|
from huggingface_hub.utils import RepositoryNotFoundError
|
||||||
|
|
||||||
|
api = HfApi()
|
||||||
|
|
||||||
|
try:
|
||||||
|
model_info: HfModelInfo = api.model_info(model_id, expand=properties)
|
||||||
|
except Exception as e:
|
||||||
|
if isinstance(e, RepositoryNotFoundError):
|
||||||
|
warnings.warn(f"{model_id} not found on Hugging Face")
|
||||||
|
model_info = None
|
||||||
|
else:
|
||||||
|
raise e
|
||||||
|
return model_info
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
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:
|
||||||
|
# Handle quant types per model size
|
||||||
|
if isinstance(quant_types, dict):
|
||||||
|
_quant_types = quant_types[size]
|
||||||
|
else:
|
||||||
|
_quant_types = quant_types
|
||||||
|
for quant_type in _quant_types:
|
||||||
|
# NOTE: models registered with org="unsloth" and QUANT_TYPE.NONE are aliases of QUANT_TYPE.UNSLOTH
|
||||||
|
_org = "unsloth" # unsloth models -- these are all quantized versions of the 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=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,
|
||||||
|
)
|
||||||
0
unsloth/utils/__init__.py
Normal file
0
unsloth/utils/__init__.py
Normal file
78
unsloth/utils/hf_hub.py
Normal file
78
unsloth/utils/hf_hub.py
Normal file
|
|
@ -0,0 +1,78 @@
|
||||||
|
from huggingface_hub import HfApi, ModelInfo
|
||||||
|
|
||||||
|
_HFAPI: HfApi = None
|
||||||
|
|
||||||
|
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.
|
||||||
|
"""
|
||||||
|
global _HFAPI
|
||||||
|
if _HFAPI is None:
|
||||||
|
_HFAPI = HfApi()
|
||||||
|
try:
|
||||||
|
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 list_models(
|
||||||
|
properties: list[str] = None,
|
||||||
|
full: bool = False,
|
||||||
|
sort: str = "downloads",
|
||||||
|
author: str = "unsloth",
|
||||||
|
search: str = None,
|
||||||
|
limit: int = 10,
|
||||||
|
) -> list[ModelInfo]:
|
||||||
|
"""
|
||||||
|
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.
|
||||||
|
sort: str = The sort order.
|
||||||
|
author: str = The author of the model.
|
||||||
|
search: str = The search query for filtering models.
|
||||||
|
|
||||||
|
"""
|
||||||
|
global _HFAPI
|
||||||
|
if _HFAPI is None:
|
||||||
|
_HFAPI = HfApi()
|
||||||
|
if full:
|
||||||
|
properties = None
|
||||||
|
|
||||||
|
models: list[ModelInfo] = _HFAPI.list_models(
|
||||||
|
author=author,
|
||||||
|
search=search,
|
||||||
|
sort=sort,
|
||||||
|
limit=limit,
|
||||||
|
expand=properties,
|
||||||
|
full=full,
|
||||||
|
)
|
||||||
|
return models
|
||||||
Loading…
Add table
Add a link
Reference in a new issue