964362818e
Do not cache transient repository-list failures as empty weight sets. Deduplicate normalized file paths without lowering reported file sizes. Record whitelisted OOM evidence and bounded discovery diagnostics.
635 lines
26 KiB
Python
635 lines
26 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, 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()
|
|
self._popular_page_lock = threading.Lock()
|
|
self._popular_page_cursor = 0
|
|
self.last_listing_diagnostics: list[dict[str, Any]] = []
|
|
|
|
@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,
|
|
sort: str = "last_modified",
|
|
rotate_pages: bool = False,
|
|
) -> list[HFModelSummary]:
|
|
if not pipeline_tags:
|
|
return []
|
|
|
|
deduped: dict[str, HFModelSummary] = {}
|
|
self.last_listing_diagnostics = []
|
|
del read_concurrency
|
|
|
|
per_tag_limit = limit if len(pipeline_tags) <= 1 else max(10, (limit + len(pipeline_tags) - 1) // len(pipeline_tags))
|
|
start_page = 1
|
|
if rotate_pages and sort == "downloads":
|
|
window_pages = max(1, min(60, (per_tag_limit + 49) // 50))
|
|
window_count = max(1, 60 // window_pages)
|
|
with self._popular_page_lock:
|
|
start_page = 1 + (self._popular_page_cursor % window_count) * window_pages
|
|
self._popular_page_cursor += 1
|
|
print(
|
|
f"[modelscope] popular_window=start_page:{start_page} pages:{window_pages}",
|
|
flush=True,
|
|
)
|
|
|
|
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,
|
|
sort=sort,
|
|
start_page=start_page,
|
|
diagnostics=self.last_listing_diagnostics,
|
|
):
|
|
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,
|
|
sort: str = "last_modified",
|
|
start_page: int = 1,
|
|
diagnostics: list[dict[str, Any]] | None = 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] = []
|
|
page_count = (max_items + page_size - 1) // page_size
|
|
diagnostic = {
|
|
"pipelineTag": pipeline_tag, "sort": sort,
|
|
"requestedLimit": limit, "effectiveLimit": max_items,
|
|
"pageSize": page_size, "startPage": start_page,
|
|
"pagesRead": 0, "cachedPages": 0, "rawRowsScanned": 0,
|
|
"acceptedModels": 0, "stopReason": "page_budget",
|
|
"catalogExhausted": False,
|
|
}
|
|
if diagnostics is not None:
|
|
diagnostics.append(diagnostic)
|
|
for page_number in range(start_page, min(61, start_page + page_count)):
|
|
cache_key = (sort, "*" 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])
|
|
diagnostic["cachedPages"] += 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": sort,
|
|
}
|
|
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
|
|
diagnostic["stopReason"] = "unsupported_filter_retry"
|
|
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,
|
|
sort=sort,
|
|
start_page=start_page,
|
|
diagnostics=diagnostics,
|
|
)
|
|
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,
|
|
)
|
|
diagnostic["stopReason"] = "http_error"
|
|
diagnostic["httpStatus"] = exc.status_code
|
|
break
|
|
|
|
items = self._extract_models(payload)
|
|
self._remember_lru(
|
|
self._model_page_cache,
|
|
cache_key,
|
|
(time.monotonic(), list(items)),
|
|
self._page_cache_limit,
|
|
)
|
|
diagnostic["pagesRead"] += 1
|
|
diagnostic["rawRowsScanned"] += len(items)
|
|
if not items:
|
|
diagnostic["stopReason"] = "empty_page"
|
|
diagnostic["catalogExhausted"] = True
|
|
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)
|
|
diagnostic["acceptedModels"] += 1
|
|
if len(models) >= max_items:
|
|
diagnostic["stopReason"] = "result_limit"
|
|
return models
|
|
if len(items) < page_size:
|
|
diagnostic["stopReason"] = "short_page"
|
|
diagnostic["catalogExhausted"] = True
|
|
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)
|
|
# A failed lookup is not an empty repository. Let the caller defer
|
|
# inspection and retry next cycle instead of caching missing
|
|
# weights as a permanent framework-incompatibility decision.
|
|
raise
|
|
|
|
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(file_sizes.get(path, 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(set(file_paths)),
|
|
file_sizes=file_sizes,
|
|
gguf_files=sorted(set(gguf_files)),
|
|
selected_gguf=PurePosixPath(selected_gguf).name if selected_gguf else None,
|
|
weight_files=sorted(set(vllm_weight_files)),
|
|
onnx_files=sorted(set(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
|