File size: 4,341 Bytes
c60ec61 | 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 | """Generate Middle-earth names from any trained checkpoint.
Works for all four architectures: it reads the checkpoint's architecture folder,
rebuilds the model, and samples with the matching BPE tokeniser. Like the worker
in train_one.py it runs one architecture per process, so the flat module names in
the architecture folders never collide.
Run: python src/generate.py qwen3 512 # 20 names, temperature 0.8
python src/generate.py gemma4 char --count 40 # the CharTokenizer baseline
python src/generate.py deepseek3 512 --temperature 0.7 --novel-only
Lower temperature -> safer, more familiar names. Higher -> more varied/inventive.
"""
from __future__ import annotations
import argparse
import importlib
import sys
from pathlib import Path
import torch
ROOT = Path(__file__).resolve().parents[1]
REPO = ROOT / "single_letter_transformers"
CKPT_DIR = ROOT / "checkpoints"
NAMES_FILE = ROOT / "data" / "middle_earth_names.txt"
# architecture -> (folder, model class name).
ARCHITECTURES = {
"qwen3": ("qwen3", "TinyQwen"),
"qwen3_5": ("qwen3_5", "TinyQwen35"),
"gemma4": ("gemma4", "TinyGemma"),
"deepseek3": ("deepseek3", "TinyDeepSeek"),
}
def load_model_and_tokenizer(arch: str, tok_spec: str):
"""Rebuild the trained model and its tokeniser (char or BPE) from the checkpoint."""
folder, class_name = ARCHITECTURES[arch]
sys.path.insert(0, str(REPO / folder))
# The architecture folder must be importable *before* torch.load, because the
# pickled cfg is that folder's config.ModelConfig.
model_class = getattr(importlib.import_module("model"), class_name)
name = "char" if tok_spec == "char" else f"bpe{tok_spec}"
checkpoint_path = CKPT_DIR / f"{arch}_{name}.pt"
if not checkpoint_path.exists():
sys.exit(f"No checkpoint at {checkpoint_path.relative_to(ROOT)} -- train it first.")
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
model = model_class(ckpt["cfg"])
model.load_state_dict(ckpt["model"])
model.eval()
if ckpt.get("tokenizer_kind") == "char":
# Rebuild the CharTokenizer from the exact character list it was trained on.
char_tokenizer = importlib.import_module("tokenizer").CharTokenizer
tokenizer = char_tokenizer(ckpt["chars"])
else:
from bpe_tokenizer import BpeTokenizer
tokenizer = BpeTokenizer.from_file(ROOT / ckpt["tokenizer"])
return model, tokenizer
@torch.no_grad()
def generate_names(model, tokenizer, count: int, temperature: float) -> list[str]:
"""Sample `count` names, each starting from the newline (start-of-name) token."""
start = torch.full((count, 1), tokenizer.newline_id, dtype=torch.long)
out = model.generate(start, max_new_tokens=model.cfg.max_seq_len,
temperature=temperature, top_k=None, eos_id=tokenizer.eos_id)
names = []
for row in out.tolist():
# Drop the leading newline, then keep everything up to the next newline.
name = tokenizer.decode(row[1:]).split("\n")[0]
if name:
names.append(name)
return names
def main() -> None:
parser = argparse.ArgumentParser(description="Sample names from a trained checkpoint.")
parser.add_argument("arch", choices=sorted(ARCHITECTURES))
parser.add_argument("tokenizer", choices=("char", "256", "512"),
help="'char' for the CharTokenizer baseline, or a BPE vocab size")
parser.add_argument("--count", type=int, default=20)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--seed", type=int, default=None, help="fix the RNG for reproducible samples")
parser.add_argument("--novel-only", action="store_true",
help="only show names that are NOT already in the training corpus")
args = parser.parse_args()
if args.seed is not None:
torch.manual_seed(args.seed)
model, tokenizer = load_model_and_tokenizer(args.arch, args.tokenizer)
names = generate_names(model, tokenizer, args.count, args.temperature)
if args.novel_only:
known = set(NAMES_FILE.read_text(encoding="utf-8").split("\n"))
names = [n for n in names if n not in known]
for name in names:
print(name)
if __name__ == "__main__":
main()
|