splitbit_llm/config.py

"""Configuration and hardware detection for SplitBit LLM."""

from __future__ import annotations

import logging
import os
import platform
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)

try:
    import psutil
except ImportError:
    psutil = None  # type: ignore[assignment]


class HardwareTier(str, Enum):
    MOBILE = "mobile"
    MINIMAL = "minimal"
    LIGHT = "light"
    STANDARD = "standard"
    FULL = "full"
    MAXIMUM = "maximum"
    DATACENTER = "datacenter"
    SUPERCOMPUTER = "supercomputer"


@dataclass
class ModelConfig:
    """Model architecture config — auto-adjusted per hardware tier."""
    n_layers: int = 3
    n_heads: int = 4
    d_model: int = 256
    d_ff: int = 512
    vocab_size: int = 4096
    max_seq_len: int = 512
    dropout: float = 0.0

    @property
    def d_head(self) -> int:
        return self.d_model // self.n_heads

    @property
    def param_count(self) -> int:
        """Rough parameter count estimate."""
        emb = self.vocab_size * self.d_model
        attn = 4 * self.d_model * self.d_model * self.n_layers
        ffn = 2 * self.d_model * self.d_ff * self.n_layers
        ln = 2 * self.d_model * self.n_layers
        return emb + attn + ffn + ln


@dataclass
class QuantConfig:
    """SplitBit weight quantization config."""
    format: str = "q4_k_m"
    bpw: float = 4.0

    @staticmethod
    def for_tier(tier: HardwareTier) -> "QuantConfig":
        formats = {
            HardwareTier.MOBILE: ("ternary", 1.585),
            HardwareTier.MINIMAL: ("q2_k", 2.0),
            HardwareTier.LIGHT: ("q3_k_s", 3.0),
            HardwareTier.STANDARD: ("q4_k_m", 4.0),
            HardwareTier.FULL: ("q5_k_m", 5.0),
            HardwareTier.MAXIMUM: ("q8_0", 8.0),
            HardwareTier.DATACENTER: ("fp8_e4m3", 8.0),
            HardwareTier.SUPERCOMPUTER: ("fp16", 16.0),
        }
        fmt, bpw = formats.get(tier, ("q4_k_m", 4.0))
        return QuantConfig(format=fmt, bpw=bpw)


@dataclass
class StorageConfig:
    """Storage paths and limits."""
    data_dir: str = os.path.expanduser("~/.splitbit-llm")
    skill_storage_mb: int = 200
    max_skills: int = 2000

    @staticmethod
    def for_tier(tier: HardwareTier) -> "StorageConfig":
        limits = {
            HardwareTier.MOBILE: (50, 500),
            HardwareTier.MINIMAL: (200, 2000),
            HardwareTier.LIGHT: (1000, 10000),
            HardwareTier.STANDARD: (5000, 50000),
            HardwareTier.FULL: (20000, 200000),
            HardwareTier.MAXIMUM: (50000, 500000),
            HardwareTier.DATACENTER: (200000, 2000000),
            HardwareTier.SUPERCOMPUTER: (1000000, 10000000),
        }
        mb, skills = limits.get(tier, (200, 2000))
        return StorageConfig(skill_storage_mb=mb, max_skills=skills)


def detect_hardware() -> HardwareTier:
    """Detect hardware tier based on available RAM, CPU, and GPU."""
    ram_mb = 0
    cpu_cores = os.cpu_count() or 1

    if psutil:
        ram_mb = psutil.virtual_memory().available // (1024 * 1024)
        cpu_cores = psutil.cpu_count() or cpu_cores

    gpu_vram_mb = 0
    try:
        import torch
        if torch.cuda.is_available():
            gpu_vram_mb = torch.cuda.get_device_properties(0).total_memory // (1024 * 1024)
    except ImportError:
        pass

    is_mobile = platform.machine().startswith(("arm", "aarch")) and ram_mb < 2048

    if is_mobile or ram_mb < 2048:
        return HardwareTier.MOBILE
    if ram_mb < 4096:
        return HardwareTier.MINIMAL
    if ram_mb < 8192:
        return HardwareTier.LIGHT
    if ram_mb < 16384 or gpu_vram_mb < 6144:
        return HardwareTier.STANDARD
    if ram_mb < 32768 or gpu_vram_mb < 12288:
        return HardwareTier.FULL
    if ram_mb < 65536:
        return HardwareTier.MAXIMUM
    if ram_mb < 262144:
        return HardwareTier.DATACENTER
    return HardwareTier.SUPERCOMPUTER


def get_model_config(tier: HardwareTier) -> ModelConfig:
    """Get optimal model config for a hardware tier."""
    configs = {
        HardwareTier.MOBILE: ModelConfig(n_layers=2, n_heads=4, d_model=128, d_ff=256, vocab_size=2048, max_seq_len=256),
        HardwareTier.MINIMAL: ModelConfig(n_layers=3, n_heads=4, d_model=256, d_ff=512, vocab_size=4096, max_seq_len=512),
        HardwareTier.LIGHT: ModelConfig(n_layers=4, n_heads=8, d_model=384, d_ff=768, vocab_size=8192, max_seq_len=1024),
        HardwareTier.STANDARD: ModelConfig(n_layers=6, n_heads=8, d_model=512, d_ff=1024, vocab_size=16384, max_seq_len=2048),
        HardwareTier.FULL: ModelConfig(n_layers=8, n_heads=16, d_model=768, d_ff=2048, vocab_size=32000, max_seq_len=4096),
        HardwareTier.MAXIMUM: ModelConfig(n_layers=12, n_heads=16, d_model=1024, d_ff=3072, vocab_size=64000, max_seq_len=8192),
        HardwareTier.DATACENTER: ModelConfig(n_layers=24, n_heads=32, d_model=2048, d_ff=6144, vocab_size=128000, max_seq_len=16384),
        HardwareTier.SUPERCOMPUTER: ModelConfig(n_layers=48, n_heads=64, d_model=4096, d_ff=12288, vocab_size=128000, max_seq_len=32768),
    }
    return configs.get(tier, configs[HardwareTier.MINIMAL])


@dataclass
class Settings:
    """Global settings for SplitBit LLM."""
    tier: HardwareTier = field(default_factory=detect_hardware)
    model: ModelConfig = field(default_factory=lambda: get_model_config(detect_hardware()))
    quant: QuantConfig = field(default_factory=lambda: QuantConfig.for_tier(detect_hardware()))
    storage: StorageConfig = field(default_factory=lambda: StorageConfig.for_tier(detect_hardware()))

    @staticmethod
    def from_env() -> "Settings":
        tier = detect_hardware()
        return Settings(
            tier=tier,
            model=get_model_config(tier),
            quant=QuantConfig.for_tier(tier),
            storage=StorageConfig.for_tier(tier),
        )