| """ |
| Self-contained loader for this checkpoint. |
| |
| Reconstructs the model architecture and tokenizer entirely from |
| config.json (no hardcoded hyperparameters here), loads the trained |
| weights, and generates a sample -- proving this folder is sufficient on |
| its own, with no dependency on the original training repo. |
| |
| Usage: |
| python load_model.py |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from bigram import BigramLanguageModel |
|
|
| HERE = Path(__file__).resolve().parent |
|
|
|
|
| def _load_weights(path: Path) -> dict[str, torch.Tensor]: |
| if path.suffix == ".safetensors": |
| from safetensors.torch import load_file |
|
|
| return load_file(path) |
| return torch.load(path, map_location="cpu", weights_only=True) |
|
|
|
|
| def load_model_and_tokenizer(dir_path: Path = HERE): |
| config = json.loads((dir_path / "config.json").read_text()) |
|
|
| if config["tokenizer_type"] == "char": |
| from tokenizer import CharTokenizer as TokenizerClass |
| elif config["tokenizer_type"] == "bpe": |
| from bpe_tokenizer import BPETokenizer as TokenizerClass |
| else: |
| raise ValueError(f"Unknown tokenizer_type: {config['tokenizer_type']!r}") |
| tokenizer = TokenizerClass.load(dir_path / config["tokenizer_file"]) |
|
|
| model = BigramLanguageModel( |
| vocab_size=config["vocab_size"], |
| block_size=config["block_size"], |
| n_embd=config["n_embd"], |
| n_head=config["n_head"], |
| n_layer=config["n_layer"], |
| ) |
| state_dict = _load_weights(dir_path / config["weights_file"]) |
| model.load_state_dict(state_dict) |
| model.eval() |
|
|
| return model, tokenizer, config |
|
|
|
|
| def main() -> None: |
| model, tokenizer, config = load_model_and_tokenizer() |
|
|
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"Loaded model: {n_params:,} parameters (config.json says {config['parameter_count']:,})") |
| assert n_params == config["parameter_count"], "Reconstructed model doesn't match config!" |
|
|
| context = torch.zeros((1, 1), dtype=torch.long) |
| print("\n----- generated sample -----") |
| sample = tokenizer.decode(model.generate(context, max_new_tokens=400)[0].tolist()) |
| print(sample) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
