Spaces:
Sleeping
Sleeping
File size: 4,049 Bytes
4a6405d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """
OpenMind Utility Functions.
Common helpers used across the project.
"""
import os
import sys
import json
import random
import logging
from pathlib import Path
from typing import Optional
import numpy as np
import torch
def setup_logging(
level: str = "INFO",
log_file: Optional[str] = None,
name: str = "openmind",
) -> logging.Logger:
"""Configure logging for the project."""
logger = logging.getLogger(name)
logger.setLevel(getattr(logging, level.upper()))
formatter = logging.Formatter(
"%(asctime)s | %(levelname)s | %(name)s | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
# Console handler
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
# File handler (optional)
if log_file:
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
return logger
def set_seed(seed: int = 42):
"""Set random seed for reproducibility across all libraries."""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def count_parameters(model: torch.nn.Module) -> dict:
"""Count model parameters by trainability."""
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
return {
"total": total,
"trainable": trainable,
"frozen": total - trainable,
"total_mb": total * 4 / (1024 * 1024), # Assuming float32
}
def get_device(prefer_gpu: bool = True) -> torch.device:
"""Get the best available device."""
if prefer_gpu and torch.cuda.is_available():
return torch.device("cuda")
elif prefer_gpu and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def format_number(n: int) -> str:
"""Format large numbers with suffixes (e.g., 125M, 1.3B)."""
if n >= 1_000_000_000:
return f"{n / 1_000_000_000:.1f}B"
elif n >= 1_000_000:
return f"{n / 1_000_000:.0f}M"
elif n >= 1_000:
return f"{n / 1_000:.0f}K"
return str(n)
def load_config(config_path: str) -> dict:
"""Load a YAML configuration file."""
import yaml
with open(config_path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
def save_json(data: dict, path: str, indent: int = 2):
"""Save dictionary as JSON file."""
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=indent, default=str)
def load_json(path: str) -> dict:
"""Load JSON file as dictionary."""
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def estimate_model_memory(config) -> dict:
"""Estimate memory requirements for a model configuration."""
# Parameters (approximate)
embed_params = config.vocab_size * config.dim
attn_params = config.n_layers * (
config.dim * config.dim * 3 # Q, K, V
+ config.dim * config.dim # O
)
ffn_params = config.n_layers * (
config.dim * config.intermediate_dim * 3 # gate, up, down
)
norm_params = config.n_layers * config.dim * 2 + config.dim # per-layer + final
total_params = embed_params + attn_params + ffn_params + norm_params
# Memory estimates (bytes)
fp32_mem = total_params * 4
fp16_mem = total_params * 2
training_mem = fp16_mem * 4 # Rough: params + grads + optimizer states
return {
"parameters": total_params,
"parameters_formatted": format_number(total_params),
"fp32_gb": fp32_mem / (1024**3),
"fp16_gb": fp16_mem / (1024**3),
"training_estimate_gb": training_mem / (1024**3),
}
|