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