File size: 5,374 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""
Interactive streaming text generation REPL for Teensy.
Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""
import os
import sys
import argparse
import codecs
import secrets
from contextlib import nullcontext
import torch
import tiktoken
from model import TeensyConfig, TeensyLM, adapt_nanogpt_weights
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 load_teensy(args):
args.device = resolve_device(args.device)
device_type = 'cuda' if 'cuda' in args.device else 'mps' if 'mps' in args.device else 'cpu'
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')
checkpoint = torch.load(ckpt_path, map_location=args.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(args.device)
enc = tiktoken.get_encoding("gpt2")
encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
token_bytes = enc.decode_single_token_bytes
eos_token_id = enc.eot_token
return model, encode, token_bytes, ctx, eos_token_id
def stream_generate(prompt, model, encode, token_bytes, ctx, args, eos_token_id):
ids = encode(prompt)
x = torch.tensor(ids, dtype=torch.long, device=args.device)[None, ...]
utf8_decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
yield prompt
with torch.no_grad():
with ctx:
for token_id in model.generate_stream(
x, args.max_new_tokens,
temperature=args.temperature,
top_k=args.top_k,
top_p=args.top_p,
eos_token_id=eos_token_id,
):
chunk = utf8_decoder.decode(token_bytes(token_id), final=False)
if chunk:
yield chunk
tail = utf8_decoder.decode(b"", final=True)
if tail:
yield tail
def main():
parser = argparse.ArgumentParser(description='Interactive Teensy text generation')
parser.add_argument('--out_dir', default='checkpoints', help='Directory containing model checkpoint')
parser.add_argument('--temperature', type=float, default=0.8, help='Sampling temperature')
parser.add_argument('--max_new_tokens', type=int, default=500, help='Number of tokens to generate')
parser.add_argument('--top_k', type=int, default=200, help='Top-k sampling parameter')
parser.add_argument('--top_p', type=float, default=0.9, help='Nucleus (top-p) sampling parameter')
parser.add_argument('--device', default='cpu', help='Device to run on (cpu/cuda/mps; cpu recommended)')
parser.add_argument('--dtype', default='float16', help='Data type for model')
parser.add_argument('--seed', type=int, default=None, help='Random seed (omit for non-deterministic sampling)')
args = parser.parse_args()
seed = args.seed if args.seed is not None else secrets.randbelow(2**32)
torch.manual_seed(seed)
if args.device == 'cuda' and torch.cuda.is_available():
torch.cuda.manual_seed(seed)
print(f"Loading model from {args.out_dir}...")
model, encode, token_bytes, ctx, eos_token_id = load_teensy(args)
print(f"Model loaded! Running on {args.device}")
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"number of parameters: {n_params/1e6:.2f}M")
print("No meta.pkl found, assuming GPT-2 encodings...")
print("\nEnter your prompts (Ctrl+C or type 'exit' to quit):")
while True:
try:
prompt = input("\n> ")
if prompt.lower() == 'exit':
break
sys.stdout.write("\n")
sys.stdout.flush()
for chunk in stream_generate(prompt, model, encode, token_bytes, ctx, args, eos_token_id):
sys.stdout.write(chunk)
sys.stdout.flush()
sys.stdout.write("\n")
sys.stdout.flush()
except KeyboardInterrupt:
sys.stdout.write("\nExiting...\n")
sys.stdout.flush()
break
except Exception as e:
print(f"Error: {e}")
if __name__ == "__main__":
main()
|