""" 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")