repl_hf.py

#!/usr/bin/env python3
"""
Interactive streaming text generation REPL for the exported Teensy model.

Loads the HuggingFace-compatible safetensors export (model.safetensors +
config.json) instead of the full PyTorch checkpoint.

Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""

import os
import sys
import json
import argparse
import codecs
import secrets
from contextlib import nullcontext

import torch
import tiktoken
from safetensors.torch import load_file

from model import TeensyConfig, TeensyLM


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)

    with open(os.path.join(args.export_dir, 'config.json')) as f:
        cfg = TeensyConfig(**json.load(f))

    model = TeensyLM(cfg)
    model.load_state_dict(load_file(os.path.join(args.export_dir, 'model.safetensors')))
    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 streaming REPL for the exported Teensy model')
    parser.add_argument('--export_dir', default='exported', help='Directory containing config.json and model.safetensors')
    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='float32', 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 exported model from {args.export_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("Using 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()