File size: 4,689 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 | #!/usr/bin/env python3
"""
Run the HuggingFace-exported Teensy model (safetensors + config.json).
This script does NOT depend on transformers; it loads the exported weights
back into the native TeensyLM architecture defined in model.py.
Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""
import os
import json
import argparse
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_model(export_dir, device, dtype):
device_type = 'cuda' if 'cuda' in device else 'mps' if 'mps' in device else 'cpu'
with open(os.path.join(export_dir, 'config.json')) as f:
cfg = TeensyConfig(**json.load(f))
model = TeensyLM(cfg)
state_dict = load_file(os.path.join(export_dir, 'model.safetensors'))
model.load_state_dict(state_dict)
model.eval()
model.to(device)
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type in ['cpu', 'mps'] else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
enc = tiktoken.get_encoding("gpt2")
encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
decode = lambda l: enc.decode(l)
eos_token_id = enc.eot_token
return model, encode, decode, ctx, eos_token_id
def generate(prompt, model, encode, decode, ctx, args, eos_token_id):
ids = encode(prompt)
x = torch.tensor(ids, dtype=torch.long, device=args.device)[None, ...]
with torch.no_grad():
with ctx:
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,
)
return decode(y[0].tolist())
def main():
parser = argparse.ArgumentParser(description='Run exported Teensy safetensors model')
parser.add_argument('--export_dir', default='exported', help='Directory containing config.json and model.safetensors')
parser.add_argument('--prompt', default='Hello', help='Input prompt')
parser.add_argument('--num_samples', type=int, default=1, help='Number of samples')
parser.add_argument('--temperature', type=float, default=0.8, help='Sampling temperature')
parser.add_argument('--max_new_tokens', type=int, default=200, 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()
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)
print(f"Loading exported model from {args.export_dir}...")
model, encode, decode, ctx, eos_token_id = load_model(args.export_dir, args.device, args.dtype)
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")
for sample_idx in range(args.num_samples):
if args.num_samples > 1:
print(f"\n--- sample {sample_idx + 1} ---")
print(generate(args.prompt, model, encode, decode, ctx, args, eos_token_id))
if __name__ == '__main__':
main()
|