583 lines
23 KiB
Python
583 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import threading
|
|
import time
|
|
from collections import OrderedDict
|
|
from dataclasses import replace
|
|
from datetime import datetime, timezone
|
|
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 = 5,
|
|
page_interval_seconds: float | None = None,
|
|
page_cache_ttl_seconds: float | None = None,
|
|
) -> 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,
|
|
backoff_seconds=2.0,
|
|
)
|
|
self.legacy_http_client = legacy_http_client or JsonHttpClient(
|
|
base_url=base_url,
|
|
default_headers=headers,
|
|
timeout=timeout,
|
|
retries=retries,
|
|
)
|
|
self._detail_cache_limit = max(
|
|
32,
|
|
min(1000, int(os.getenv("MODELSCOPE_DETAIL_CACHE_MAX_MODELS", "128"))),
|
|
)
|
|
self._page_cache_limit = max(
|
|
16,
|
|
min(500, int(os.getenv("MODELSCOPE_PAGE_CACHE_MAX_PAGES", "96"))),
|
|
)
|
|
self._repo_tree_cache: OrderedDict[str, list[dict[str, Any]]] = OrderedDict()
|
|
self._repo_tree_lock = threading.Lock()
|
|
self._model_config_cache: OrderedDict[str, tuple[dict[str, Any], str | None]] = OrderedDict()
|
|
self._model_config_lock = threading.Lock()
|
|
self._model_last_modified_cache: OrderedDict[str, datetime | None] = OrderedDict()
|
|
self._model_last_modified_lock = threading.Lock()
|
|
self._model_page_cache: OrderedDict[
|
|
tuple[str, int, int], tuple[float, list[dict[str, Any]]]
|
|
] = OrderedDict()
|
|
self._model_page_cache_ttl = max(
|
|
0.0,
|
|
float(
|
|
page_cache_ttl_seconds
|
|
if page_cache_ttl_seconds is not None
|
|
else os.getenv("MODELSCOPE_PAGE_CACHE_TTL_SECONDS", "900")
|
|
),
|
|
)
|
|
self._page_interval_seconds = max(
|
|
0.0,
|
|
float(
|
|
page_interval_seconds
|
|
if page_interval_seconds is not None
|
|
else os.getenv("MODELSCOPE_PAGE_INTERVAL_SECONDS", "0.25")
|
|
),
|
|
)
|
|
self._last_model_page_request_at = 0.0
|
|
self._unsupported_task_filters: dict[str, float] = {}
|
|
self._model_card_cache: OrderedDict[str, dict[str, Any]] = OrderedDict()
|
|
self._model_card_lock = threading.Lock()
|
|
|
|
@staticmethod
|
|
def _remember_lru(cache: OrderedDict, key: Any, value: Any, limit: int) -> None:
|
|
cache[key] = value
|
|
cache.move_to_end(key)
|
|
while len(cache) > limit:
|
|
cache.popitem(last=False)
|
|
|
|
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
|
|
|
|
per_tag_limit = limit if len(pipeline_tags) <= 1 else max(10, (limit + len(pipeline_tags) - 1) // len(pipeline_tags))
|
|
|
|
for pipeline_tag in pipeline_tags:
|
|
for model in self._query_recent_models(
|
|
pipeline_tag=pipeline_tag,
|
|
limit=per_tag_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[: max(1, int(limit))]
|
|
|
|
def _query_recent_models(
|
|
self,
|
|
*,
|
|
pipeline_tag: str,
|
|
limit: int,
|
|
min_downloads: int,
|
|
updated_after=None,
|
|
) -> list[HFModelSummary]:
|
|
# ModelScope OpenAPI rejects values above 50 with InputParameterError.
|
|
page_size = min(max(1, limit), 50)
|
|
max_items = min(max(1, limit), 3000)
|
|
task_tag = MODELSCOPE_TASK_TAGS.get(pipeline_tag, pipeline_tag)
|
|
filter_disabled = self._unsupported_task_filters.get(task_tag, 0.0) > time.monotonic()
|
|
models: list[HFModelSummary] = []
|
|
for page_number in range(1, (max_items + page_size - 1) // page_size + 1):
|
|
cache_key = ("*" if filter_disabled else task_tag, page_number, page_size)
|
|
cached = self._model_page_cache.get(cache_key)
|
|
if cached is not None and time.monotonic() - cached[0] < self._model_page_cache_ttl:
|
|
self._model_page_cache.move_to_end(cache_key)
|
|
items = list(cached[1])
|
|
else:
|
|
elapsed = time.monotonic() - self._last_model_page_request_at
|
|
if self._last_model_page_request_at > 0 and elapsed < self._page_interval_seconds:
|
|
time.sleep(self._page_interval_seconds - elapsed)
|
|
try:
|
|
query = {
|
|
"page_number": page_number,
|
|
"page_size": page_size,
|
|
"sort": "last_modified",
|
|
}
|
|
if not filter_disabled:
|
|
query["filter.task"] = task_tag
|
|
payload = self.http_client.request_json(
|
|
"GET",
|
|
"/models",
|
|
query=query,
|
|
)
|
|
self._last_model_page_request_at = time.monotonic()
|
|
except HttpJsonError as exc:
|
|
self._last_model_page_request_at = time.monotonic()
|
|
if exc.status_code == 400 and not filter_disabled:
|
|
self._unsupported_task_filters[task_tag] = time.monotonic() + 86_400
|
|
print(
|
|
f"[modelscope] task_filter_unsupported task={task_tag} "
|
|
"fallback=unfiltered ttl=86400s",
|
|
flush=True,
|
|
)
|
|
return self._query_recent_models(
|
|
pipeline_tag=pipeline_tag,
|
|
limit=limit,
|
|
min_downloads=min_downloads,
|
|
updated_after=updated_after,
|
|
)
|
|
print(
|
|
f"[modelscope] list_models_error task={task_tag} page={page_number} "
|
|
f"partial_models={len(models)} retry_next_cycle=true error={exc}",
|
|
flush=True,
|
|
)
|
|
break
|
|
|
|
items = self._extract_models(payload)
|
|
self._remember_lru(
|
|
self._model_page_cache,
|
|
cache_key,
|
|
(time.monotonic(), list(items)),
|
|
self._page_cache_limit,
|
|
)
|
|
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)
|
|
inspection = inspect_repo_tree(model.repo_id, entries)
|
|
if not inspection.has_root_config:
|
|
return replace(inspection, published_size_bytes=model.file_size)
|
|
model_config, config_error = self.get_model_config(model.repo_id)
|
|
model_card_metadata = self.get_model_card_metadata(model.repo_id)
|
|
return ModelInspection(
|
|
repo_id=inspection.repo_id,
|
|
file_paths=inspection.file_paths,
|
|
file_sizes=inspection.file_sizes,
|
|
gguf_files=inspection.gguf_files,
|
|
selected_gguf=inspection.selected_gguf,
|
|
weight_files=inspection.weight_files,
|
|
onnx_files=inspection.onnx_files,
|
|
model_config=model_config,
|
|
config_fetch_error=config_error,
|
|
model_card_metadata=model_card_metadata,
|
|
published_size_bytes=model.file_size,
|
|
)
|
|
|
|
def get_model_card_metadata(self, repo_id: str) -> dict[str, Any]:
|
|
with self._model_card_lock:
|
|
cached = self._model_card_cache.get(repo_id)
|
|
if cached is not None:
|
|
self._model_card_cache.move_to_end(repo_id)
|
|
if cached is not None:
|
|
return dict(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}")
|
|
data = payload.get("Data") if isinstance(payload, dict) else None
|
|
if not isinstance(data, dict) and isinstance(payload, dict):
|
|
data = payload.get("data")
|
|
readme = ""
|
|
if isinstance(data, dict):
|
|
readme = str(data.get("ReadMe") or data.get("readme") or data.get("README") or "")
|
|
result = parse_model_card_front_matter(readme)
|
|
except Exception:
|
|
result = {}
|
|
with self._model_card_lock:
|
|
self._remember_lru(
|
|
self._model_card_cache,
|
|
repo_id,
|
|
dict(result),
|
|
self._detail_cache_limit,
|
|
)
|
|
return result
|
|
|
|
def get_model_config(self, repo_id: str) -> tuple[dict[str, Any], str | None]:
|
|
with self._model_config_lock:
|
|
cached = self._model_config_cache.get(repo_id)
|
|
if cached is not None:
|
|
self._model_config_cache.move_to_end(repo_id)
|
|
if cached is not None:
|
|
return dict(cached[0]), cached[1]
|
|
|
|
encoded_repo_id = "/".join(quote(part, safe="") for part in repo_id.split("/"))
|
|
try:
|
|
payload = self.legacy_http_client.request_json(
|
|
"GET",
|
|
f"/models/{encoded_repo_id}/resolve/master/config.json",
|
|
)
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("config.json did not contain a JSON object")
|
|
result = (dict(payload), None)
|
|
except Exception as exc:
|
|
result = ({}, f"{type(exc).__name__}: {exc}")
|
|
|
|
with self._model_config_lock:
|
|
self._remember_lru(
|
|
self._model_config_cache,
|
|
repo_id,
|
|
result,
|
|
self._detail_cache_limit,
|
|
)
|
|
return dict(result[0]), result[1]
|
|
|
|
def get_model_last_modified(self, repo_id: str) -> datetime | None:
|
|
with self._model_last_modified_lock:
|
|
if repo_id in self._model_last_modified_cache:
|
|
result = self._model_last_modified_cache[repo_id]
|
|
self._model_last_modified_cache.move_to_end(repo_id)
|
|
return result
|
|
|
|
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}",
|
|
)
|
|
except HttpJsonError as exc:
|
|
print(f"[modelscope] model_metadata_error repo={repo_id} error={exc}", flush=True)
|
|
payload = None
|
|
|
|
data = payload.get("Data") if isinstance(payload, dict) else None
|
|
if not isinstance(data, dict) and isinstance(payload, dict):
|
|
data = payload.get("data")
|
|
result: datetime | None = None
|
|
if isinstance(data, dict):
|
|
raw_value = (
|
|
data.get("LastUpdatedTime")
|
|
or data.get("lastUpdatedTime")
|
|
or data.get("last_modified")
|
|
or data.get("updated_at")
|
|
)
|
|
if isinstance(raw_value, (int, float)):
|
|
timestamp = float(raw_value)
|
|
if timestamp > 10_000_000_000:
|
|
timestamp /= 1000.0
|
|
try:
|
|
result = datetime.fromtimestamp(timestamp, tz=timezone.utc)
|
|
except (OverflowError, OSError, ValueError):
|
|
result = None
|
|
else:
|
|
result = parse_datetime(raw_value)
|
|
|
|
with self._model_last_modified_lock:
|
|
self._remember_lru(
|
|
self._model_last_modified_cache,
|
|
repo_id,
|
|
result,
|
|
self._detail_cache_limit,
|
|
)
|
|
return result
|
|
|
|
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:
|
|
self._repo_tree_cache.move_to_end(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._remember_lru(
|
|
self._repo_tree_cache,
|
|
repo_id,
|
|
list(entries),
|
|
self._detail_cache_limit,
|
|
)
|
|
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")),
|
|
params=_optional_int(item.get("params") or item.get("Params") or item.get("parameter_count")),
|
|
file_size=_optional_int(item.get("file_size") or item.get("FileSize") or item.get("size")),
|
|
tags=_string_tuple(item.get("tags") or item.get("Tags")),
|
|
tasks=_string_tuple(item.get("tasks") or item.get("Tasks")),
|
|
license=str(item.get("license") or item.get("License") or "").strip() or None,
|
|
gated=bool(item.get("gated") or item.get("Gated")),
|
|
private=bool(item.get("private") or item.get("Private")),
|
|
likes=int(item.get("likes") or item.get("Likes") or 0),
|
|
)
|
|
|
|
|
|
def _optional_int(value: Any) -> int | None:
|
|
try:
|
|
parsed = int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
return parsed if parsed >= 0 else None
|
|
|
|
|
|
def _string_tuple(value: Any) -> tuple[str, ...]:
|
|
if isinstance(value, str):
|
|
values = [part.strip() for part in value.split(",")]
|
|
elif isinstance(value, (list, tuple, set)):
|
|
values = [str(part).strip() for part in value]
|
|
else:
|
|
values = []
|
|
return tuple(dict.fromkeys(part for part in values if part))
|
|
|
|
|
|
def parse_model_card_front_matter(readme: str) -> dict[str, Any]:
|
|
text = str(readme or "")[:65_536]
|
|
if not text.startswith("---"):
|
|
return {}
|
|
match = re.match(r"^---\s*\n(.*?)\n---(?:\s*\n|$)", text, flags=re.DOTALL)
|
|
if match is None:
|
|
return {}
|
|
front_matter = match.group(1)
|
|
try:
|
|
import yaml # type: ignore
|
|
|
|
parsed = yaml.safe_load(front_matter)
|
|
data = parsed if isinstance(parsed, dict) else {}
|
|
except (ImportError, ValueError, TypeError):
|
|
data = {}
|
|
current_list: str | None = None
|
|
for raw_line in front_matter.splitlines():
|
|
if re.match(r"^\s+-\s+", raw_line) and current_list:
|
|
value = re.sub(r"^\s+-\s+", "", raw_line).strip().strip('"\'')
|
|
data.setdefault(current_list, []).append(value)
|
|
continue
|
|
if ":" not in raw_line or raw_line[:1].isspace():
|
|
continue
|
|
key, value = raw_line.split(":", 1)
|
|
key = key.strip()
|
|
value = value.strip().strip('"\'')
|
|
if not value:
|
|
data[key] = []
|
|
current_list = key
|
|
else:
|
|
data[key] = value
|
|
current_list = None
|
|
allowed = ("base_model", "base_model_relation", "frameworks", "tasks", "new_version")
|
|
return {key: data[key] for key in allowed if key in data}
|
|
|
|
|
|
def inspect_repo_tree(repo_id: str, entries: list[dict[str, Any]]) -> ModelInspection:
|
|
file_paths: list[str] = []
|
|
file_sizes: dict[str, int] = {}
|
|
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
|
|
path = str(path)
|
|
if path.startswith("./"):
|
|
path = path[2:]
|
|
path = path.lstrip("/")
|
|
file_paths.append(path)
|
|
size_value: Any = None
|
|
size_present = False
|
|
for size_key in ("Size", "size"):
|
|
if size_key in entry:
|
|
size_value = entry[size_key]
|
|
size_present = True
|
|
break
|
|
if size_present:
|
|
try:
|
|
file_sizes[path] = max(0, int(size_value))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
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),
|
|
file_sizes=file_sizes,
|
|
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
|