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