Model/utils/common.py

"""
Frox AI Morph 1.1 — Utilities
Small, dependency-light helpers shared across training, inference, and
scripts. Morph 1.0 had no equivalent module — every script rolled its
own seeding / device-detection logic, which is how subtle
non-reproducibility bugs creep in.
"""
from __future__ import annotations

import os
import random
import sys
import time
from contextlib import contextmanager
from typing import Optional

import numpy as np
import torch


# ── Reproducibility ───────────────────────────────────────────────

def set_seed(seed: int = 1337, deterministic: bool = False):
    """
    Seed every RNG Morph touches: Python, NumPy, PyTorch (CPU + all CUDA
    devices). `deterministic=True` additionally forces cuDNN into
    deterministic mode — slower, but bit-exact reruns for debugging.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    else:
        torch.backends.cudnn.benchmark = True

    print(f"✓ Seed set to {seed} (deterministic={deterministic})")


# ── Device detection ───────────────────────────────────────────────

def get_device(prefer: Optional[str] = None) -> torch.device:
    """
    Pick the best available device.
    prefer: force "cuda" | "mps" | "cpu" if given and available.
    """
    if prefer:
        if prefer == "cuda" and torch.cuda.is_available():
            return torch.device("cuda")
        if prefer == "mps" and torch.backends.mps.is_available():
            return torch.device("mps")
        if prefer == "cpu":
            return torch.device("cpu")
        print(f"⚠ Requested device '{prefer}' unavailable, auto-detecting instead")

    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def describe_device(device: torch.device) -> str:
    """Human-readable device description for logs."""
    if device.type == "cuda":
        idx = device.index or 0
        name = torch.cuda.get_device_name(idx)
        total_gb = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 3)
        cc = torch.cuda.get_device_capability(idx)
        return f"{name} ({total_gb:.1f}GB, compute capability {cc[0]}.{cc[1]})"
    if device.type == "mps":
        return "Apple Silicon (MPS)"
    return "CPU"


def recommended_dtype(device: torch.device) -> torch.dtype:
    """bfloat16 on Ampere+ (A100/H100/RTX 30xx+), float16 on older CUDA, float32 on CPU/MPS."""
    if device.type != "cuda":
        return torch.float32
    cc = torch.cuda.get_device_capability(device)
    return torch.bfloat16 if cc[0] >= 8 else torch.float16


# ── Environment detection ─────────────────────────────────────────

def detect_environment() -> str:
    """Return 'kaggle' | 'colab' | 'local' for auto-configuring save paths."""
    if os.path.exists("/kaggle/working"):
        return "kaggle"
    if os.path.exists("/content"):
        return "colab"
    return "local"


def require_checkpoint_dir(path: str) -> "Path":
    """
    Confirm `path` is a checkpoint directory containing config.json before
    any loader tries to open it, and fail with an actionable message instead
    of a bare FileNotFoundError pointing at the config.json open() call.

    The most common cause: a phase (pretrain/sft) was skipped or never
    finished, so the checkpoint directory --from-checkpoint / --model points
    at was never written.
    """
    from pathlib import Path
    p = Path(path)

    if not p.exists():
        parent = p.parent
        siblings = sorted(d.name for d in parent.iterdir() if d.is_dir()) if parent.exists() else []
        hint = (
            f"\n  Checkpoints found in {parent}: {', '.join(siblings)}"
            if siblings else
            f"\n  {parent} doesn't exist yet or is empty — no training phase has saved a checkpoint there."
        )
        raise FileNotFoundError(
            f"Checkpoint directory not found: {p}{hint}\n"
            f"  If you skipped an earlier phase (e.g. pretrain), either run that phase first "
            f"or drop --from-checkpoint / --model to start from fresh weights."
        )

    if not (p / "config.json").exists():
        contents = sorted(f.name for f in p.iterdir())
        raise FileNotFoundError(
            f"{p} exists but has no config.json (found: {', '.join(contents) or 'nothing'}).\n"
            f"  This usually means the save that was supposed to write here didn't complete — "
            f"check the log for the run that was meant to produce this checkpoint."
        )

    return p


# ── Timing ────────────────────────────────────────────────────────

@contextmanager
def timer(label: str = "block"):
    """Context manager that prints elapsed wall-clock time on exit."""
    t0 = time.perf_counter()
    yield
    elapsed = time.perf_counter() - t0
    print(f"⏱  {label}: {elapsed:.2f}s")


# ── Logging setup ─────────────────────────────────────────────────

def setup_logging(level: str = "INFO"):
    """Configure structlog if available, else fall back to stdlib logging."""
    try:
        import structlog
        structlog.configure(
            processors=[
                structlog.processors.TimeStamper(fmt="iso"),
                structlog.processors.add_log_level,
                structlog.dev.ConsoleRenderer(),
            ],
        )
        return structlog.get_logger()
    except ImportError:
        import logging
        logging.basicConfig(
            level=getattr(logging, level),
            format="%(asctime)s [%(levelname)s] %(message)s",
            stream=sys.stdout,
        )
        return logging.getLogger("morph")


# ── Misc ──────────────────────────────────────────────────────────

def human_readable_bytes(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if abs(n) < 1024.0:
            return f"{n:.1f}{unit}"
        n /= 1024.0
    return f"{n:.1f}PB"


# ── Model family loading ───────────────────────────────────────────

FAMILY_TIERS = ("nano", "mini", "classic", "pro", "code")
_LEGACY_SCALE_ALIASES = {"1.5b": "mini", "1b": "mini", "3b": "classic", "8b": "code", "7b": "code"}


def load_family_config(name: str):
    """
    Load exactly one Morph model-family tier config, on demand.

    Uses importlib so requesting "nano" never imports mini/classic/
    pro/code's modules — each tier stays independently loadable, which
    matters if you're shipping only one tier's files in a deployment.
    """
    import importlib

    key = name.strip().lower()
    key = _LEGACY_SCALE_ALIASES.get(key, key)

    if key not in FAMILY_TIERS:
        raise ValueError(
            f"Unknown Morph family '{name}'. Choose from: {', '.join(FAMILY_TIERS)} "
            f"(legacy aliases also work: {', '.join(_LEGACY_SCALE_ALIASES)})"
        )

    module = importlib.import_module(f"config.family.{key}")
    return module.get_config(), module


def print_banner(version: str = "1.1.0"):
    print(r"""
    ███████╗██████╗  ██████╗ ██╗  ██╗
    ██╔════╝██╔══██╗██╔═══██╗╚██╗██╔╝
    █████╗  ██████╔╝██║   ██║ ╚███╔╝
    ██╔══╝  ██╔══██╗██║   ██║ ██╔██╗
    ██║     ██║  ██║╚██████╔╝██╔╝ ██╗
    ╚═╝     ╚═╝  ╚═╝ ╚═════╝ ╚═╝  ╚═╝
    """)
    print(f"    Frox AI — Morph {version}\n")