85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
from huggingface_hub import ModelInfo as HfModelInfo
|
|
from unsloth.model_registry import (
|
|
ModelInfo,
|
|
get_llama_models,
|
|
get_llama_vision_models,
|
|
get_phi_instruct_models,
|
|
get_phi_models,
|
|
get_qwen_models,
|
|
get_qwen_vl_models,
|
|
)
|
|
from unsloth.utils.hf_hub import get_model_info
|
|
|
|
MODEL_NAMES = [
|
|
"llama",
|
|
"llama_vision",
|
|
"qwen",
|
|
"qwen_vl",
|
|
"phi",
|
|
"phi_instruct",
|
|
]
|
|
REGISTERED_MODELS = [
|
|
get_llama_models(),
|
|
get_llama_vision_models(),
|
|
get_qwen_models(),
|
|
get_qwen_vl_models(),
|
|
get_phi_models(),
|
|
get_phi_instruct_models(),
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class ModelTestParam:
|
|
name: str
|
|
models: dict[str, ModelInfo]
|
|
|
|
|
|
def _test_model_uploaded(model_ids: list[str]):
|
|
missing_models = []
|
|
for _id in model_ids:
|
|
model_info: HfModelInfo = get_model_info(_id)
|
|
if not model_info:
|
|
missing_models.append(_id)
|
|
|
|
return missing_models
|
|
|
|
|
|
TestParams = [
|
|
ModelTestParam(name, models)
|
|
for name, models in zip(MODEL_NAMES, REGISTERED_MODELS)
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_test_param", TestParams, ids=lambda param: param.name
|
|
)
|
|
def test_model_uploaded(model_test_param: ModelTestParam):
|
|
missing_models = _test_model_uploaded(model_test_param.models)
|
|
assert not missing_models, (
|
|
f"{model_test_param.name} missing following models: {missing_models}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for method in [
|
|
get_llama_models,
|
|
get_llama_vision_models,
|
|
get_qwen_models,
|
|
get_qwen_vl_models,
|
|
get_phi_models,
|
|
get_phi_instruct_models,
|
|
]:
|
|
models = method()
|
|
model_name = next(iter(models.values())).base_name
|
|
print(f"{model_name}: {len(models)} registered")
|
|
for model_info in models.values():
|
|
print(f" {model_info.model_path}")
|
|
missing_models = test_model_uploaded(list(models.keys()))
|
|
|
|
if missing_models:
|
|
print("--------------------------------")
|
|
print(f"Missing models: {missing_models}")
|
|
print("--------------------------------")
|