230 lines
7.7 KiB
Python
230 lines
7.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
|
|
import os
|
||
|
|
import threading
|
||
|
|
from pathlib import PurePosixPath
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from common import parse_datetime
|
||
|
|
from http_json import JsonHttpClient
|
||
|
|
from models import HFModelSummary, ModelInspection
|
||
|
|
|
||
|
|
|
||
|
|
GGUF_PRIORITY = ("q4_0.gguf", "q8_0.gguf", "fp16.gguf")
|
||
|
|
VLLM_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth")
|
||
|
|
ONNX_WEIGHT_SUFFIXES = (".onnx",)
|
||
|
|
|
||
|
|
|
||
|
|
class HuggingFaceDiscovery:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
base_url: str = "https://huggingface.co",
|
||
|
|
http_client: JsonHttpClient | None = None,
|
||
|
|
timeout: int = 30,
|
||
|
|
retries: int = 2,
|
||
|
|
) -> None:
|
||
|
|
hf_token = os.getenv("HF_TOKEN")
|
||
|
|
headers = {"User-Agent": "modelhub-submmit-cli/0.1"}
|
||
|
|
if hf_token:
|
||
|
|
headers["Authorization"] = f"Bearer {hf_token}"
|
||
|
|
self.http_client = http_client or JsonHttpClient(
|
||
|
|
base_url=base_url,
|
||
|
|
default_headers=headers,
|
||
|
|
timeout=timeout,
|
||
|
|
retries=retries,
|
||
|
|
)
|
||
|
|
self._repo_tree_cache: dict[str, list[dict[str, Any]]] = {}
|
||
|
|
self._repo_tree_lock = threading.Lock()
|
||
|
|
|
||
|
|
def list_recent_models(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
pipeline_tags: list[str],
|
||
|
|
limit: int,
|
||
|
|
min_downloads: int,
|
||
|
|
updated_after=None,
|
||
|
|
read_concurrency: int = 1,
|
||
|
|
) -> list[HFModelSummary]:
|
||
|
|
if not pipeline_tags:
|
||
|
|
return []
|
||
|
|
|
||
|
|
deduped: dict[str, HFModelSummary] = {}
|
||
|
|
max_workers = min(len(pipeline_tags), max(1, read_concurrency))
|
||
|
|
|
||
|
|
if max_workers <= 1:
|
||
|
|
tag_results = [
|
||
|
|
self._query_recent_models(
|
||
|
|
pipeline_tag=pipeline_tag,
|
||
|
|
limit=limit,
|
||
|
|
min_downloads=min_downloads,
|
||
|
|
updated_after=updated_after,
|
||
|
|
)
|
||
|
|
for pipeline_tag in pipeline_tags
|
||
|
|
]
|
||
|
|
else:
|
||
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
|
|
futures = {
|
||
|
|
executor.submit(
|
||
|
|
self._query_recent_models,
|
||
|
|
pipeline_tag=pipeline_tag,
|
||
|
|
limit=limit,
|
||
|
|
min_downloads=min_downloads,
|
||
|
|
updated_after=updated_after,
|
||
|
|
): pipeline_tag
|
||
|
|
for pipeline_tag in pipeline_tags
|
||
|
|
}
|
||
|
|
tag_results = []
|
||
|
|
for future in as_completed(futures):
|
||
|
|
tag_results.append(future.result())
|
||
|
|
|
||
|
|
for models in tag_results:
|
||
|
|
for model in models:
|
||
|
|
current = deduped.get(model.repo_id)
|
||
|
|
if current is None or (model.last_modified or parse_datetime("1970-01-01")) > (
|
||
|
|
current.last_modified or parse_datetime("1970-01-01")
|
||
|
|
):
|
||
|
|
deduped[model.repo_id] = model
|
||
|
|
models = list(deduped.values())
|
||
|
|
models.sort(key=lambda item: item.last_modified or parse_datetime("1970-01-01"), reverse=True)
|
||
|
|
return models
|
||
|
|
|
||
|
|
def _query_recent_models(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
pipeline_tag: str,
|
||
|
|
limit: int,
|
||
|
|
min_downloads: int,
|
||
|
|
updated_after=None,
|
||
|
|
) -> list[HFModelSummary]:
|
||
|
|
payload = self.http_client.request_json(
|
||
|
|
"GET",
|
||
|
|
"/api/models",
|
||
|
|
query={
|
||
|
|
"pipeline_tag": pipeline_tag,
|
||
|
|
"sort": "lastModified",
|
||
|
|
"direction": "-1",
|
||
|
|
"limit": limit,
|
||
|
|
"full": "true",
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
models: list[HFModelSummary] = []
|
||
|
|
for item in payload or []:
|
||
|
|
model = self._parse_model(item, min_downloads=min_downloads, updated_after=updated_after)
|
||
|
|
if model is not None:
|
||
|
|
models.append(model)
|
||
|
|
return models
|
||
|
|
|
||
|
|
def list_recent_text_generation_models(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
limit: int,
|
||
|
|
min_downloads: int,
|
||
|
|
updated_after=None,
|
||
|
|
) -> list[HFModelSummary]:
|
||
|
|
return self.list_recent_models(
|
||
|
|
pipeline_tags=["text-generation"],
|
||
|
|
limit=limit,
|
||
|
|
min_downloads=min_downloads,
|
||
|
|
updated_after=updated_after,
|
||
|
|
)
|
||
|
|
|
||
|
|
def inspect_model(self, model: HFModelSummary) -> ModelInspection:
|
||
|
|
entries = self.list_repo_tree(model.repo_id)
|
||
|
|
return inspect_repo_tree(model.repo_id, entries)
|
||
|
|
|
||
|
|
def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]:
|
||
|
|
with self._repo_tree_lock:
|
||
|
|
cached = self._repo_tree_cache.get(repo_id)
|
||
|
|
if cached is not None:
|
||
|
|
return list(cached)
|
||
|
|
|
||
|
|
payload = self.http_client.request_json(
|
||
|
|
"GET",
|
||
|
|
f"/api/models/{repo_id}/tree/main",
|
||
|
|
query={"recursive": "1"},
|
||
|
|
)
|
||
|
|
|
||
|
|
entries: list[dict[str, Any]]
|
||
|
|
if isinstance(payload, list):
|
||
|
|
entries = payload
|
||
|
|
elif isinstance(payload, dict):
|
||
|
|
for key in ("items", "tree", "siblings"):
|
||
|
|
value = payload.get(key)
|
||
|
|
if isinstance(value, list):
|
||
|
|
entries = value
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
entries = []
|
||
|
|
else:
|
||
|
|
entries = []
|
||
|
|
|
||
|
|
with self._repo_tree_lock:
|
||
|
|
self._repo_tree_cache[repo_id] = list(entries)
|
||
|
|
return list(entries)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _parse_model(item: dict[str, Any], *, min_downloads: int, updated_after=None) -> HFModelSummary | None:
|
||
|
|
repo_id = item.get("id") or item.get("modelId")
|
||
|
|
if not repo_id:
|
||
|
|
return None
|
||
|
|
downloads = int(item.get("downloads") or 0)
|
||
|
|
if downloads < min_downloads:
|
||
|
|
return None
|
||
|
|
pipeline_tag = item.get("pipeline_tag") or item.get("pipelineTag")
|
||
|
|
last_modified = parse_datetime(item.get("lastModified") or item.get("last_modified"))
|
||
|
|
if updated_after and last_modified and last_modified < updated_after:
|
||
|
|
return None
|
||
|
|
return HFModelSummary(
|
||
|
|
repo_id=repo_id,
|
||
|
|
downloads=downloads,
|
||
|
|
last_modified=last_modified,
|
||
|
|
pipeline_tag=pipeline_tag,
|
||
|
|
created_at=parse_datetime(item.get("createdAt") or item.get("created_at")),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def inspect_repo_tree(repo_id: str, entries: list[dict[str, Any]]) -> ModelInspection:
|
||
|
|
file_paths: list[str] = []
|
||
|
|
gguf_files: list[str] = []
|
||
|
|
vllm_weight_files: list[str] = []
|
||
|
|
onnx_files: list[str] = []
|
||
|
|
|
||
|
|
for entry in entries:
|
||
|
|
path = entry.get("path") or entry.get("rfilename") or entry.get("name")
|
||
|
|
if not path:
|
||
|
|
continue
|
||
|
|
entry_type = (entry.get("type") or "").lower()
|
||
|
|
if entry_type in {"directory", "dir", "folder"}:
|
||
|
|
continue
|
||
|
|
file_paths.append(path)
|
||
|
|
filename = PurePosixPath(path).name.lower()
|
||
|
|
if any(filename.endswith(suffix) for suffix in GGUF_PRIORITY):
|
||
|
|
gguf_files.append(path)
|
||
|
|
if filename.endswith(VLLM_WEIGHT_SUFFIXES):
|
||
|
|
vllm_weight_files.append(path)
|
||
|
|
if filename.endswith(ONNX_WEIGHT_SUFFIXES):
|
||
|
|
onnx_files.append(path)
|
||
|
|
|
||
|
|
selected_gguf = choose_best_gguf(gguf_files)
|
||
|
|
return ModelInspection(
|
||
|
|
repo_id=repo_id,
|
||
|
|
file_paths=sorted(file_paths),
|
||
|
|
gguf_files=sorted(gguf_files),
|
||
|
|
selected_gguf=PurePosixPath(selected_gguf).name if selected_gguf else None,
|
||
|
|
weight_files=sorted(vllm_weight_files),
|
||
|
|
onnx_files=sorted(onnx_files),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def choose_best_gguf(paths: list[str]) -> str | None:
|
||
|
|
if not paths:
|
||
|
|
return None
|
||
|
|
ranked_paths = sorted(paths)
|
||
|
|
for suffix in GGUF_PRIORITY:
|
||
|
|
for path in ranked_paths:
|
||
|
|
if PurePosixPath(path).name.lower().endswith(suffix):
|
||
|
|
return path
|
||
|
|
return None
|