File size: 4,256 Bytes
081a81c | 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 110 | #!/usr/bin/env python3
"""
Generate text samples from a trained Teensy checkpoint.
Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""
import os
import argparse
import secrets
from contextlib import nullcontext
import torch
import tiktoken
from model import TeensyConfig, TeensyLM, adapt_nanogpt_weights
def load_model(checkpoint_path, device):
checkpoint = torch.load(checkpoint_path, map_location=device)
cfg = TeensyConfig(**checkpoint['model_args'])
model = TeensyLM(cfg)
state_dict = checkpoint['model']
unwanted_prefix = '_orig_mod.'
for k, v in list(state_dict.items()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
state_dict = adapt_nanogpt_weights(state_dict)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
return model
def resolve_device(requested):
"""Pick an available device and warn about MPS quality issues."""
if requested == "cuda":
if torch.cuda.is_available():
return "cuda"
fallback = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"CUDA not available; falling back to {fallback}.")
return fallback
if requested == "mps":
if torch.backends.mps.is_available():
print("WARNING: MPS inference for this model can produce degraded output")
print(" (garbled text and stray <|endoftext|> tokens). CPU is recommended.")
return "mps"
print("MPS not available; falling back to cpu.")
return "cpu"
return "cpu"
def main():
parser = argparse.ArgumentParser(description="Generate text from a Teensy checkpoint")
parser.add_argument("--out_dir", default="checkpoints", help="Checkpoint directory")
parser.add_argument("--start", default="\n", help="Prompt text")
parser.add_argument("--num_samples", type=int, default=10, help="Number of samples")
parser.add_argument("--max_new_tokens", type=int, default=500, help="Tokens per sample")
parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
parser.add_argument("--top_k", type=int, default=50, help="Top-k sampling")
parser.add_argument("--top_p", type=float, default=0.9, help="Nucleus (top-p) sampling")
parser.add_argument("--device", default="cpu", help="Device (cpu/cuda/mps; cpu recommended)")
parser.add_argument("--dtype", default="float16", help="Torch dtype")
parser.add_argument("--compile", action="store_true", help="torch.compile the model")
parser.add_argument("--seed", type=int, default=None, help="Random seed")
args = parser.parse_args()
args.device = resolve_device(args.device)
device_type = "cuda" if "cuda" in args.device else "mps" if "mps" in args.device else "cpu"
seed = args.seed if args.seed is not None else secrets.randbelow(2**32)
torch.manual_seed(seed)
if device_type == "cuda":
torch.cuda.manual_seed(seed)
ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
ctx = nullcontext() if device_type in ["cpu", "mps"] else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
ckpt_path = os.path.join(args.out_dir, "teensy-0.pt")
model = load_model(ckpt_path, args.device)
if args.compile and device_type == "cuda":
model = torch.compile(model)
enc = tiktoken.get_encoding("gpt2")
encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
decode = lambda ids: enc.decode(ids)
eos_token_id = enc.eot_token
if args.start.startswith("FILE:"):
with open(args.start[5:], "r", encoding="utf-8") as f:
args.start = f.read()
start_ids = encode(args.start)
x = torch.tensor(start_ids, dtype=torch.long, device=args.device)[None, ...]
with torch.no_grad():
with ctx:
for _ in range(args.num_samples):
y = model.generate(x, args.max_new_tokens, temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, eos_token_id=eos_token_id)
print(decode(y[0].tolist()))
print("---------------")
if __name__ == "__main__":
main()
|