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.
233 lines
7.6 KiB
Python
233 lines
7.6 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 exact_repository_size_bytes(self) -> int | None:
|
|
"""Return exact recursive on-disk size when every file has a size."""
|
|
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)
|
|
return total if total > 0 else None
|
|
return None
|
|
|
|
@property
|
|
def repository_size_bytes(self) -> int | None:
|
|
"""Use the conservative maximum of recursive and published sizes."""
|
|
exact = self.exact_repository_size_bytes
|
|
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)
|