submmit/modelhub_submmit_api/models.py

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime


@dataclass(frozen=True)
class TemplateRecord:
    template_id: str
    task_type: str
    framework: str
    target_gpu: str
    model_address_template: str
    config_params: str
    source_file: str | None = None
    source_line: int | None = None


@dataclass(frozen=True)
class HFModelSummary:
    repo_id: str
    downloads: int
    last_modified: datetime | None
    pipeline_tag: str | None
    created_at: datetime | None = None

    @property
    def model_address(self) -> str:
        return f"https://modelscope.cn/models/{self.repo_id}"


@dataclass(frozen=True)
class ModelInspection:
    repo_id: str
    file_paths: list[str] = field(default_factory=list)
    gguf_files: list[str] = field(default_factory=list)
    selected_gguf: str | None = None
    weight_files: list[str] = field(default_factory=list)
    onnx_files: list[str] = field(default_factory=list)

    @property
    def has_gguf(self) -> bool:
        return self.selected_gguf is not None

    @property
    def has_vllm_weights(self) -> bool:
        return bool(self.weight_files)

    @property
    def has_onnx_weights(self) -> bool:
        return bool(self.onnx_files)

    @property
    def has_standard_weights(self) -> bool:
        return bool(self.weight_files or self.onnx_files)


@dataclass(frozen=True)
class CandidateModel:
    repo_id: str
    model_address: str
    pipeline_tag: str | None
    modality: str
    task_type: str
    target_gpu: str
    framework: str
    template_id: str
    config_params: str
    downloads: int
    last_modified: datetime | None
    gguf_filename: str | None = None
    score: float = 0.0
    warnings: list[str] = field(default_factory=list)