diff --git a/src/modeldock/cli/commands/info.py b/src/modeldock/cli/commands/info.py index 6184bd2..a289b9a 100644 --- a/src/modeldock/cli/commands/info.py +++ b/src/modeldock/cli/commands/info.py @@ -25,9 +25,9 @@ def _echo_spec(spec: Any) -> None: if spec.variants: typer.echo("Variants:") for v in spec.variants: - size = f"{v.size_bytes} bytes" if v.size_bytes else "?" + size = str(v.size) if v.size else "?" ram = v.min_ram or "?" - typer.echo(f" - {v.tag} ({v.params or '?'}), {size}, min RAM {ram}") + typer.echo(f" - {v.tag} ({size}), min RAM {ram}") def info_cmd( diff --git a/src/modeldock/domain/model.py b/src/modeldock/domain/model.py index cec63b5..8824a7c 100644 --- a/src/modeldock/domain/model.py +++ b/src/modeldock/domain/model.py @@ -9,7 +9,7 @@ from enum import Enum from typing import TYPE_CHECKING, List, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field if TYPE_CHECKING: from modeldock.ports.registry import RegistryPort @@ -75,14 +75,46 @@ def from_value(cls, value: str) -> RuntimeBackend: raise ValueError(f"Unknown runtime backend: {value!r}") +class ModelSize(BaseModel): + """Value object representing a model's size: parameter count and disk footprint. + + Pure data — no I/O. Consolidates the loose ``params``/``size_bytes`` fields + previously carried directly on ``ModelVariant``. See issue #98. + """ + + model_config = ConfigDict(frozen=True) + + params: int # raw parameter count, e.g. 8_000_000_000 + size_bytes: int # size on disk in bytes, e.g. 4_700_000_000 + + def format_params(self) -> str: + """Return a compact parameter count, e.g. '8B' or '1.5B'.""" + if self.params >= 1_000_000_000: + value = self.params / 1_000_000_000 + suffix = "B" + else: + value = self.params / 1_000_000 + suffix = "M" + text = f"{value:.1f}".rstrip("0").rstrip(".") + return f"{text}{suffix}" + + def format_size(self) -> str: + """Return a human-readable disk size, e.g. '4.7GB'.""" + gb = self.size_bytes / (1000**3) + text = f"{gb:.1f}".rstrip("0").rstrip(".") + return f"{text}GB" + + def __str__(self) -> str: + return f"{self.format_params()} ({self.format_size()})" + + class ModelVariant(BaseModel): """A specific tag/variant of a model (e.g. llama3:8b).""" tag: str download_url: Optional[str] = None sha256: Optional[str] = None - params: Optional[str] = None - size_bytes: Optional[int] = None + size: Optional[ModelSize] = None min_ram: Optional[str] = None diff --git a/tests/unit/test_cli_json_output.py b/tests/unit/test_cli_json_output.py index 52e6243..7b6ceec 100644 --- a/tests/unit/test_cli_json_output.py +++ b/tests/unit/test_cli_json_output.py @@ -30,6 +30,7 @@ Device, ModelInfo, ModelRef, + ModelSize, ModelSpec, ModelVariant, RuntimeBackend, @@ -45,7 +46,9 @@ category=Category.CHAT, capabilities=[Capability.CHAT, Capability.TOOL_USE], default_tag="8b", - variants=[ModelVariant(tag="8b", params="8B", size_bytes=4_700_000_000)], + variants=[ + ModelVariant(tag="8b", size=ModelSize(params=8_000_000_000, size_bytes=4_700_000_000)), + ], description="Meta Llama 3", source=OLLAMA_OFFICIAL, ) diff --git a/tests/unit/test_model_size.py b/tests/unit/test_model_size.py new file mode 100644 index 0000000..39b8f62 --- /dev/null +++ b/tests/unit/test_model_size.py @@ -0,0 +1,38 @@ +"""Tests for the ModelSize value object. See issue #98.""" + +import pytest + +from modeldock.domain.model import ModelSize + + +def test_format_params_billions(): + size = ModelSize(params=8_000_000_000, size_bytes=4_700_000_000) + assert size.format_params() == "8B" + + +def test_format_params_with_decimal(): + size = ModelSize(params=1_500_000_000, size_bytes=1_000_000_000) + assert size.format_params() == "1.5B" + + +def test_format_params_millions(): + size = ModelSize(params=350_000_000, size_bytes=200_000_000) + assert size.format_params() == "350M" + + +def test_format_size_gb(): + size = ModelSize(params=8_000_000_000, size_bytes=4_700_000_000) + assert size.format_size() == "4.7GB" + + +def test_str_combines_both(): + size = ModelSize(params=8_000_000_000, size_bytes=4_700_000_000) + assert str(size) == "8B (4.7GB)" + + +def test_immutable(): + from pydantic import ValidationError + + size = ModelSize(params=8_000_000_000, size_bytes=4_700_000_000) + with pytest.raises(ValidationError): + size.params = 1 # type: ignore[misc]