File size: 2,385 Bytes
8ce1de8 | 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 | """Shared helpers: seeding, config loading/overriding, device selection, checkpoint loading."""
from __future__ import annotations
import copy
import random
from pathlib import Path
from typing import Any
import numpy as np
import torch
import yaml
PROJECT_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_CONFIG = PROJECT_ROOT / "configs" / "tiny_gpt.yaml"
def set_seed(seed: int) -> None:
"""Make runs reproducible (python, numpy, torch)."""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def get_device(prefer: str = "auto") -> torch.device:
if prefer != "auto":
return torch.device(prefer)
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
def load_config(path: str | Path | None = None) -> dict:
path = Path(path) if path else DEFAULT_CONFIG
with open(path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
def _parse_value(s: str) -> Any:
try:
return yaml.safe_load(s) # "4" -> 4, "0.1" -> 0.1, "true" -> True, "abc" -> "abc"
except Exception:
return s
def apply_overrides(cfg: dict, overrides: list[str] | None) -> dict:
"""Apply CLI overrides of the form section.key=value, e.g. model.n_layers=2."""
cfg = copy.deepcopy(cfg)
for item in overrides or []:
key, _, value = item.partition("=")
parts = key.split(".")
node = cfg
for p in parts[:-1]:
node = node.setdefault(p, {})
node[parts[-1]] = _parse_value(value)
return cfg
def resolve(path: str | Path) -> Path:
"""Resolve a path relative to the project root unless it is absolute."""
p = Path(path)
return p if p.is_absolute() else PROJECT_ROOT / p
def load_checkpoint(path: str | Path, device: str | torch.device = "cpu"):
"""Load a checkpoint written by train.py. Returns (model, tokenizer, ckpt_dict).
The checkpoint is self-contained: weights + model config + tokenizer vocabulary.
"""
from src.gpt import GPT
from src.tokenizer import CharTokenizer
ckpt = torch.load(resolve(path), map_location=device, weights_only=False)
model = GPT(**ckpt["model_config"]).to(device)
model.load_state_dict(ckpt["model_state"])
model.eval()
tok = CharTokenizer(ckpt["chars"])
return model, tok, ckpt
|