228 lines
7.5 KiB
Python
228 lines
7.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import PurePosixPath
|
|
from typing import Any
|
|
|
|
|
|
@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
|
|
params: int | None = None
|
|
file_size: int | None = None
|
|
tags: tuple[str, ...] = ()
|
|
tasks: tuple[str, ...] = ()
|
|
license: str | None = None
|
|
gated: bool = False
|
|
private: bool = False
|
|
likes: int = 0
|
|
|
|
@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)
|
|
file_sizes: dict[str, int] = field(default_factory=dict)
|
|
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)
|
|
model_config: dict[str, Any] = field(default_factory=dict)
|
|
config_fetch_error: str | None = None
|
|
model_card_metadata: dict[str, Any] = field(default_factory=dict)
|
|
published_size_bytes: int | None = None
|
|
|
|
@property
|
|
def root_file_names(self) -> set[str]:
|
|
return {
|
|
PurePosixPath(path).name.lower()
|
|
for path in self.file_paths
|
|
if len(PurePosixPath(path).parts) == 1
|
|
}
|
|
|
|
@property
|
|
def has_root_config(self) -> bool:
|
|
return "config.json" in self.root_file_names
|
|
|
|
@property
|
|
def has_root_tokenizer(self) -> bool:
|
|
names = self.root_file_names
|
|
exact_names = {
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"tokenizer.model",
|
|
"sentencepiece.bpe.model",
|
|
"sentencepiece.model",
|
|
"spiece.model",
|
|
"vocab.json",
|
|
"vocab.txt",
|
|
}
|
|
return bool(names & exact_names) or any(
|
|
name.startswith(("tokenizer_", "tokenization_")) and name.endswith(".py")
|
|
for name in names
|
|
)
|
|
|
|
@property
|
|
def has_root_standard_weights(self) -> bool:
|
|
direct = any(
|
|
len(PurePosixPath(path).parts) == 1
|
|
for path in [*self.weight_files, *self.onnx_files]
|
|
)
|
|
if direct:
|
|
return True
|
|
names = self.root_file_names
|
|
has_index = any(
|
|
name.endswith((".safetensors.index.json", ".bin.index.json"))
|
|
for name in names
|
|
)
|
|
return has_index and bool(self.weight_files or self.onnx_files)
|
|
|
|
@property
|
|
def model_type(self) -> str | None:
|
|
value = self.model_config.get("model_type")
|
|
return str(value).strip() if value not in (None, "") else None
|
|
|
|
@property
|
|
def architectures(self) -> list[str]:
|
|
value = self.model_config.get("architectures")
|
|
if isinstance(value, list):
|
|
return [str(item).strip() for item in value if str(item).strip()]
|
|
if value not in (None, ""):
|
|
return [str(value).strip()]
|
|
return []
|
|
|
|
@property
|
|
def quantization_method(self) -> str | None:
|
|
value = self.model_config.get("quantization_config")
|
|
if isinstance(value, dict):
|
|
method = value.get("quant_method") or value.get("quantization_method")
|
|
if method not in (None, ""):
|
|
return str(method).strip().lower()
|
|
return None
|
|
|
|
@property
|
|
def max_context_length(self) -> int | None:
|
|
for key in (
|
|
"max_position_embeddings",
|
|
"model_max_length",
|
|
"seq_length",
|
|
"n_positions",
|
|
"max_seq_len",
|
|
):
|
|
value = self.model_config.get(key)
|
|
try:
|
|
parsed = int(value)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
if 0 < parsed <= 10_000_000:
|
|
return parsed
|
|
return None
|
|
|
|
@property
|
|
def repository_size_bytes(self) -> int | None:
|
|
"""Return exact recursive on-disk size when every file has a size."""
|
|
exact: int | None = None
|
|
if self.file_paths and not any(path not in self.file_sizes for path in self.file_paths):
|
|
total = sum(max(0, int(self.file_sizes[path])) for path in self.file_paths)
|
|
exact = total if total > 0 else None
|
|
published = int(self.published_size_bytes or 0) or None
|
|
values = [value for value in (exact, published) if value is not None]
|
|
return max(values) if values else None
|
|
|
|
def estimated_load_bytes(self, framework: str) -> int | None:
|
|
if framework == "llamacpp":
|
|
if not self.selected_gguf:
|
|
return None
|
|
for path in self.gguf_files:
|
|
if PurePosixPath(path).name == self.selected_gguf:
|
|
size = int(self.file_sizes.get(path) or 0)
|
|
return size or None
|
|
return None
|
|
|
|
if "onnx" in framework or "sherpa" in framework:
|
|
sizes = [
|
|
int(self.file_sizes.get(path) or 0)
|
|
for path in self.onnx_files
|
|
if len(PurePosixPath(path).parts) == 1
|
|
]
|
|
total = sum(size for size in sizes if size > 0)
|
|
return total or None
|
|
|
|
# Repositories occasionally publish both .bin and .safetensors copies.
|
|
# The runtime loads one complete format, so use the smallest positive
|
|
# root-level format total instead of double-counting alternatives.
|
|
root_names = self.root_file_names
|
|
has_root_index = any(
|
|
name.endswith((".safetensors.index.json", ".bin.index.json"))
|
|
for name in root_names
|
|
)
|
|
totals: dict[str, int] = {}
|
|
for path in self.weight_files:
|
|
# A root index may legally reference shards in subdirectories. In
|
|
# that case include every shard of each format so large indexed
|
|
# checkpoints cannot evade the preflight size calculation.
|
|
if not has_root_index and len(PurePosixPath(path).parts) != 1:
|
|
continue
|
|
suffix = PurePosixPath(path).suffix.lower()
|
|
size = int(self.file_sizes.get(path) or 0)
|
|
if size > 0:
|
|
totals[suffix] = totals.get(suffix, 0) + size
|
|
positive = [value for value in totals.values() if value > 0]
|
|
return min(positive) if positive else None
|
|
|
|
@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)
|
|
preflight_metadata: dict[str, Any] = field(default_factory=dict)
|