soup/soup_cli/autopilot/analyzer.py

"""Autopilot analyzers — dataset, model, and hardware profiling."""

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Optional


@dataclass
class DatasetProfile:
    samples: int
    format: str
    avg_tokens: int
    p95_tokens: int
    quality: float


@dataclass
class ModelProfile:
    name: str
    params_b: float
    context: int
    architecture: str
    modality: str


@dataclass
class HardwareProfile:
    device: str
    gpu_name: str
    vram_gb: float
    compute_capability: float
    system_ram_gb: float


def analyze_dataset(path: str) -> DatasetProfile:
    """Inspect a dataset file and produce a profile.

    Delegates to Soup's existing loader/validator/format-detector so we don't
    duplicate heuristics.
    """
    from soup_cli.data.formats import detect_format
    from soup_cli.data.loader import load_raw_data
    from soup_cli.data.validator import extended_stats, validate_and_stats

    file_path = Path(path)
    if not file_path.exists():
        raise ValueError(f"Dataset not found: {file_path}")

    data = load_raw_data(file_path)
    if not data:
        raise ValueError("Dataset is empty")

    try:
        fmt = detect_format(data)
    except ValueError:
        fmt = "unknown"

    ext = extended_stats(data)
    stats = validate_and_stats(data, expected_format=fmt if fmt != "unknown" else None)

    # Rough quality score — 1.0 minus penalty for issues
    total_issues = sum(
        1 for key in ("empty_fields", "duplicates") if stats.get(key, 0) > 0
    )
    quality = max(0.0, 1.0 - 0.15 * total_issues - 0.05 * len(stats.get("issues", [])))
    quality = round(quality, 2)

    # token_counts is already approximate tokens; compute p95 directly.
    token_counts = sorted(ext.get("token_counts") or [])
    if token_counts:
        p95_idx = min(len(token_counts) - 1, int(len(token_counts) * 0.95))
        p95_tokens = int(token_counts[p95_idx])
    else:
        p95_tokens = int(ext["avg_tokens"])
    if p95_tokens <= 0:
        p95_tokens = max(1, int(ext["avg_tokens"]))

    return DatasetProfile(
        samples=stats["total"],
        format=fmt,
        avg_tokens=int(ext["avg_tokens"]),
        p95_tokens=p95_tokens,
        quality=quality,
    )


_MODEL_SIZE_RE = re.compile(r"(\d+(?:\.\d+)?)\s*[Bb]")


def _guess_params_from_name(name: str) -> float:
    """Extract the parameter count (in billions) from a model name."""
    match = _MODEL_SIZE_RE.search(name)
    if match:
        return float(match.group(1))
    # Conservative default
    return 7.0


def analyze_model(name: str, params_b: Optional[float] = None) -> ModelProfile:
    """Build a model profile from a HF model name.

    Uses ``utils.gpu.model_size_from_name`` when available, otherwise estimates
    from the model name. Context length defaults to 8192 unless the name
    signals otherwise (e.g. ``-128k``, ``-longctx``).
    """
    if params_b is None:
        try:
            from soup_cli.utils.gpu import model_size_from_name

            params_b = float(model_size_from_name(name)) or _guess_params_from_name(name)
        except Exception:  # noqa: BLE001
            params_b = _guess_params_from_name(name)

    lowered = name.lower()
    context = 8192
    if "128k" in lowered or "long" in lowered:
        context = 131072
    elif "32k" in lowered:
        context = 32768
    elif "qwen" in lowered or "llama-3" in lowered or "llama3" in lowered:
        context = 8192

    architecture = "dense"
    if "moe" in lowered or "a3b" in lowered or "deepseek-v3" in lowered:
        architecture = "moe"

    modality = "text"
    if "vision" in lowered or "vl-" in lowered:
        modality = "vision"
    elif "audio" in lowered:
        modality = "audio"

    return ModelProfile(
        name=name,
        params_b=float(params_b),
        context=context,
        architecture=architecture,
        modality=modality,
    )


def analyze_hardware() -> HardwareProfile:
    """Profile the current training hardware (GPU or CPU fallback)."""
    try:
        from soup_cli.utils.gpu import detect_device, get_gpu_info

        gpu_info = get_gpu_info()
        device = detect_device()
        vram_bytes = gpu_info.get("memory_total_bytes", 0) or 0
        vram_gb = vram_bytes / 1024**3
        gpu_name = gpu_info.get("name", "unknown") or "unknown"
    except Exception:  # noqa: BLE001
        device = "cpu"
        gpu_name = "none"
        vram_gb = 0.0

    # Compute capability: best-effort, falls back to 0
    compute_capability = 0.0
    try:
        import torch

        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
            major, minor = torch.cuda.get_device_capability(0)
            compute_capability = float(f"{major}.{minor}")
    except Exception:  # noqa: BLE001
        compute_capability = 0.0

    try:
        import psutil

        system_ram_gb = psutil.virtual_memory().total / 1024**3
    except Exception:  # noqa: BLE001
        system_ram_gb = 0.0

    return HardwareProfile(
        device=str(device),
        gpu_name=str(gpu_name),
        vram_gb=float(vram_gb),
        compute_capability=compute_capability,
        system_ram_gb=float(system_ram_gb),
    )