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 @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 @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.""" if not self.file_paths or any(path not in self.file_sizes for path in self.file_paths): return None total = sum(max(0, int(self.file_sizes[path])) for path in self.file_paths) return total if total > 0 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)