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