adapt for ModelHub agent platform
This commit is contained in:
CoolBoy
229
modelhub_submmit_api/hf_discovery.py
Normal file
229
modelhub_submmit_api/hf_discovery.py
Normal file
@@ -0,0 +1,229 @@
|
||||
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
|
||||
Reference in New Issue
Block a user