submmit/modelhub_submmit_api/hf_discovery.py

from __future__ import annotations

import os
import re
import threading
import time
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._repo_tree_cache: dict[str, list[dict[str, Any]]] = {}
        self._repo_tree_lock = threading.Lock()
        self._model_config_cache: dict[str, tuple[dict[str, Any], str | None]] = {}
        self._model_config_lock = threading.Lock()
        self._model_last_modified_cache: dict[str, datetime | None] = {}
        self._model_last_modified_lock = threading.Lock()
        self._model_page_cache: dict[tuple[str, int, int], tuple[float, list[dict[str, Any]]]] = {}
        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: dict[str, dict[str, Any]] = {}
        self._model_card_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

        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:
                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._model_page_cache[cache_key] = (time.monotonic(), list(items))
            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:
            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._model_card_cache[repo_id] = dict(result)
        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:
            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._model_config_cache[repo_id] = result
        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:
                return self._model_last_modified_cache[repo_id]

        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._model_last_modified_cache[repo_id] = result
        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:
            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")),
            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