codekingpro/portable-devtools

1from dataclasses import dataclass, field2from enum import Enum3from typing import Any4 5 6@dataclass(frozen=True)7class ModelSource:8    hf: str | None = None9    url: str | None = None10    _deprecated_tar_struct: bool = False11 12    @property13    def deprecated_tar_struct(self) -> bool:14        return self._deprecated_tar_struct15 16    def __post_init__(self) -> None:17        if self.hf is None and self.url is None:18            raise ValueError(19                f"At least one source should be set, current sources: hf={self.hf}, url={self.url}"20            )21 22 23@dataclass(frozen=True)24class BaseModelDescription:25    model: str26    sources: ModelSource27    model_file: str28    description: str29    license: str30    size_in_GB: float31    additional_files: list[str] = field(default_factory=list)32 33 34@dataclass(frozen=True)35class DenseModelDescription(BaseModelDescription):36    dim: int | None = None37    tasks: dict[str, Any] | None = field(default_factory=dict)38 39    def __post_init__(self) -> None:40        assert self.dim is not None, "dim is required for dense model description"41 42 43@dataclass(frozen=True)44class SparseModelDescription(BaseModelDescription):45    requires_idf: bool | None = None46    vocab_size: int | None = None47 48 49class PoolingType(str, Enum):50    CLS = "CLS"51    MEAN = "MEAN"52    DISABLED = "DISABLED"53