zsgdp/gpu/model_server.py

"""GPU model server configuration."""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(slots=True)
class GPUModelConfig:
    backend: str = "transformers"
    provider: str = "huggingface_spaces"
    space_name: str = "zeroshotGPU"
    max_batch_size: int = 4

    @classmethod
    def from_config(cls, config: dict) -> "GPUModelConfig":
        gpu = config.get("gpu", {})
        return cls(
            backend=gpu.get("backend", "transformers"),
            provider=gpu.get("provider", "huggingface_spaces"),
            space_name=gpu.get("space_name", "zeroshotGPU"),
            max_batch_size=int(gpu.get("max_batch_size", 4)),
        )