284 lines
10 KiB
Python
284 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
from pathlib import PurePosixPath
|
|
from typing import Any
|
|
from urllib.parse import quote
|
|
|
|
from common import parse_datetime
|
|
from defaults import EMBEDDED_MODELSCOPE_TOKEN
|
|
from http_json import HttpJsonError
|
|
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",)
|
|
|
|
MODELSCOPE_TASK_TAGS = {
|
|
"text-generation": "text-generation",
|
|
"image-text-to-text": "image-text-to-text",
|
|
"visual-question-answering": "visual-question-answering",
|
|
"document-question-answering": "document-question-answering",
|
|
"video-text-to-text": "video-text-to-text",
|
|
"text-to-image": "text-to-image-synthesis",
|
|
"image-to-image": "image-to-image",
|
|
"automatic-speech-recognition": "auto-speech-recognition",
|
|
"question-answering": "question-answering",
|
|
"feature-extraction": "feature-extraction",
|
|
"sentence-similarity": "sentence-similarity",
|
|
"image-classification": "image-classification",
|
|
"zero-shot-image-classification": "zero-shot-image-classification",
|
|
"text-classification": "text-classification",
|
|
"zero-shot-classification": "zero-shot-classification",
|
|
"reinforcement-learning": "reinforcement-learning",
|
|
}
|
|
|
|
|
|
class HuggingFaceDiscovery:
|
|
def __init__(
|
|
self,
|
|
base_url: str = "https://modelscope.cn",
|
|
http_client: JsonHttpClient | None = None,
|
|
legacy_http_client: JsonHttpClient | None = None,
|
|
timeout: int = 30,
|
|
retries: int = 2,
|
|
) -> None:
|
|
token = os.getenv("MODELSCOPE_API_TOKEN") or os.getenv("MODELSCOPE_TOKEN") or EMBEDDED_MODELSCOPE_TOKEN
|
|
headers = {"User-Agent": "modelhub-submmit-cli/0.1"}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
headers["Cookie"] = f"m_session_id={token}"
|
|
self.http_client = http_client or JsonHttpClient(
|
|
base_url=f"{base_url.rstrip('/')}/openapi/v1",
|
|
default_headers=headers,
|
|
timeout=timeout,
|
|
retries=retries,
|
|
)
|
|
self.legacy_http_client = legacy_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] = {}
|
|
del read_concurrency
|
|
|
|
for pipeline_tag in pipeline_tags:
|
|
for model in self._query_recent_models(
|
|
pipeline_tag=pipeline_tag,
|
|
limit=limit,
|
|
min_downloads=min_downloads,
|
|
updated_after=updated_after,
|
|
):
|
|
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]:
|
|
page_size = min(max(1, limit), 100)
|
|
max_items = min(max(1, limit), 3000)
|
|
task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag)
|
|
models: list[HFModelSummary] = []
|
|
for page_number in range(1, (max_items + page_size - 1) // page_size + 1):
|
|
try:
|
|
payload = self.http_client.request_json(
|
|
"GET",
|
|
"/models",
|
|
query={
|
|
"page_number": page_number,
|
|
"page_size": page_size,
|
|
"sort": "last_modified",
|
|
"filter.task": task_tag,
|
|
},
|
|
)
|
|
except HttpJsonError as exc:
|
|
print(f"[modelscope] list_models_error task={task_tag} page={page_number} error={exc}", flush=True)
|
|
break
|
|
|
|
items = self._extract_models(payload)
|
|
if not items:
|
|
break
|
|
for item in items:
|
|
model = self._parse_model(
|
|
item,
|
|
fallback_pipeline_tag=pipeline_tag,
|
|
min_downloads=min_downloads,
|
|
updated_after=updated_after,
|
|
)
|
|
if model is not None:
|
|
models.append(model)
|
|
if len(models) >= max_items:
|
|
return models
|
|
if len(items) < page_size:
|
|
break
|
|
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)
|
|
|
|
encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
|
|
try:
|
|
payload = self.legacy_http_client.request_json(
|
|
"GET",
|
|
f"/api/v1/models/{encoded_repo_id}/repo/files",
|
|
query={"Revision": "master", "Recursive": "true"},
|
|
)
|
|
except HttpJsonError as exc:
|
|
print(f"[modelscope] repo_tree_error repo={repo_id} error={exc}", flush=True)
|
|
payload = None
|
|
|
|
entries = self._extract_files(payload)
|
|
|
|
with self._repo_tree_lock:
|
|
self._repo_tree_cache[repo_id] = list(entries)
|
|
return list(entries)
|
|
|
|
@staticmethod
|
|
def _extract_models(payload: Any) -> list[dict[str, Any]]:
|
|
data = payload.get("data") if isinstance(payload, dict) else payload
|
|
if isinstance(data, dict):
|
|
for key in ("models", "Models", "items", "list", "data", "results"):
|
|
value = data.get(key)
|
|
if isinstance(value, list):
|
|
return [item for item in value if isinstance(item, dict)]
|
|
if isinstance(data, list):
|
|
return [item for item in data if isinstance(item, dict)]
|
|
return []
|
|
|
|
@staticmethod
|
|
def _extract_files(payload: Any) -> list[dict[str, Any]]:
|
|
data = payload.get("Data") if isinstance(payload, dict) else payload
|
|
if isinstance(data, dict):
|
|
for key in ("Files", "files", "items", "tree"):
|
|
value = data.get(key)
|
|
if isinstance(value, list):
|
|
return [item for item in value if isinstance(item, dict)]
|
|
if isinstance(data, list):
|
|
return [item for item in data if isinstance(item, dict)]
|
|
return []
|
|
|
|
@staticmethod
|
|
def _parse_model(
|
|
item: dict[str, Any],
|
|
*,
|
|
fallback_pipeline_tag: str,
|
|
min_downloads: int,
|
|
updated_after=None,
|
|
) -> HFModelSummary | None:
|
|
repo_id = item.get("id") or item.get("model_id") or item.get("modelId")
|
|
if not repo_id:
|
|
owner = item.get("owner") or item.get("Owner") or item.get("Path")
|
|
name = item.get("name") or item.get("Name")
|
|
repo_id = f"{owner}/{name}" if owner and name else None
|
|
if not repo_id:
|
|
return None
|
|
downloads = int(item.get("downloads") or item.get("Downloads") or 0)
|
|
if downloads < min_downloads:
|
|
return None
|
|
pipeline_tag = fallback_pipeline_tag
|
|
last_modified = parse_datetime(item.get("last_modified") or item.get("UpdatedAt") or item.get("LastUpdatedTime"))
|
|
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("created_at") or item.get("CreatedAt")),
|
|
)
|
|
|
|
|
|
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("Path") or entry.get("rfilename") or entry.get("name") or entry.get("Name")
|
|
if not path:
|
|
continue
|
|
entry_type = (entry.get("type") or 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
|