submmit/modelhub_submmit_api/hf_discovery.py

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