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