File size: 2,216 Bytes
436c9f7 50aa7f9 436c9f7 50aa7f9 436c9f7 50aa7f9 436c9f7 50aa7f9 436c9f7 | 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 | """
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()
|