| """ |
| Local inference script for MicroGPT. |
| Loads a trained checkpoint and generates text. |
| |
| Usage: |
| # Interactive mode (type prompts) |
| python inference.py |
| |
| # Generate from a prompt |
| python inference.py --prompt "Once upon a time" |
| |
| # Use a specific checkpoint |
| python inference.py --checkpoint runs_out/micro_gpt_ckpt_step_100000.pt --prompt "Hello" |
| """ |
|
|
| import os |
| import sys |
| import argparse |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from models import get_model |
|
|
|
|
| def load_model_and_meta(checkpoint_path: str, device: str = 'cpu'): |
| """Load a trained model checkpoint and return model + metadata.""" |
| print(f'Loading checkpoint: {checkpoint_path}') |
| ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) |
|
|
| config = ckpt['config'] |
| meta = ckpt['meta'] |
| step = ckpt['step'] |
| best_val_loss = ckpt.get('best_val_loss', 'N/A') |
|
|
| print(f' Architecture : {config["arch"]}') |
| print(f' Dataset : {config["dataset"]}') |
| print(f' Vocab size : {meta["vocab_size"]}') |
| print(f' Block size : {config["block_size"]}') |
| print(f' Layers/Heads : {config["n_layer"]}/{config["n_head"]}') |
| print(f' Embed dim : {config["n_embd"]}') |
| print(f' Parameters : {ckpt["n_params"]:,}') |
| print(f' Training step: {step}') |
| print(f' Best val loss: {best_val_loss}') |
|
|
| |
| ModelClass = get_model(config['arch']) |
| model = ModelClass( |
| vocab_size=meta['vocab_size'], |
| block_size=config['block_size'], |
| n_layer=config['n_layer'], |
| n_head=config['n_head'], |
| n_embd=config['n_embd'], |
| dropout=config.get('dropout', 0.1), |
| ) |
| model.load_state_dict(ckpt['model']) |
| model.to(device) |
| model.eval() |
|
|
| return model, meta, config |
|
|
|
|
| def encode_text(text: str, stoi: dict) -> torch.Tensor: |
| """Encode a string into token indices using the char-level tokenizer.""" |
| indices = [stoi.get(c, 0) for c in text] |
| return torch.tensor([indices], dtype=torch.long) |
|
|
|
|
| def decode_tokens(tokens: list, itos: dict) -> str: |
| """Decode token indices back into a string.""" |
| return ''.join(itos.get(i, '?') for i in tokens) |
|
|
|
|
| @torch.no_grad() |
| def generate_text(model, prompt: str, meta: dict, max_new_tokens: int = 256, |
| temperature: float = 0.9, top_k: int = 40, device: str = 'cpu'): |
| """Generate text continuation from a prompt.""" |
| stoi = meta['stoi'] |
| itos = meta['itos'] |
|
|
| |
| prompt_ids = encode_text(prompt, stoi) |
| prompt_ids = prompt_ids.to(device) |
|
|
| |
| output_ids = model.generate( |
| prompt_ids, |
| max_new_tokens=max_new_tokens, |
| temperature=temperature, |
| top_k=top_k, |
| ) |
|
|
| |
| full_text = decode_tokens(output_ids[0].tolist(), itos) |
| return full_text |
|
|
|
|
| def interactive_mode(model, meta, device: str): |
| """Interactive chat-like generation.""" |
| print('\n' + '=' * 60) |
| print(' MicroGPT Interactive Mode') |
| print(' Type your prompt and press Enter.') |
| print(' Type "quit", "exit", or "q" to stop.') |
| print('=' * 60) |
|
|
| while True: |
| try: |
| prompt = input('\nPrompt: ').strip() |
| except (EOFError, KeyboardInterrupt): |
| print() |
| break |
|
|
| if prompt.lower() in ('quit', 'exit', 'q'): |
| break |
|
|
| if not prompt: |
| continue |
|
|
| print('\nGenerating...\n') |
| result = generate_text(model, prompt, meta, device=device) |
| print(result) |
| print() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='MicroGPT Local Inference') |
| parser.add_argument('--checkpoint', type=str, |
| default='runs_out/micro_gpt_ckpt_step_100000.pt', |
| help='Path to model checkpoint') |
| parser.add_argument('--prompt', type=str, default=None, |
| help='Text prompt for generation') |
| parser.add_argument('--max_new_tokens', type=int, default=256, |
| help='Maximum tokens to generate') |
| parser.add_argument('--temperature', type=float, default=0.9, |
| help='Sampling temperature (higher = more random)') |
| parser.add_argument('--top_k', type=int, default=40, |
| help='Top-k sampling threshold') |
| parser.add_argument('--device', type=str, default='cpu', |
| help='Device to run on (cpu, cuda, mps)') |
| args = parser.parse_args() |
|
|
| |
| if args.device == 'cuda' and not torch.cuda.is_available(): |
| print('CUDA not available, falling back to CPU') |
| args.device = 'cpu' |
| elif args.device == 'mps' and not hasattr(torch.backends, 'mps') or not torch.backends.mps.is_available(): |
| print('MPS not available, falling back to CPU') |
| args.device = 'cpu' |
|
|
| |
| model, meta, config = load_model_and_meta(args.checkpoint, args.device) |
|
|
| if args.prompt: |
| |
| print(f'\nPrompt: {args.prompt}') |
| print(f'Generating (temp={args.temperature}, top_k={args.top_k}, max_tokens={args.max_new_tokens})...\n') |
| result = generate_text( |
| model, args.prompt, meta, |
| max_new_tokens=args.max_new_tokens, |
| temperature=args.temperature, |
| top_k=args.top_k, |
| device=args.device, |
| ) |
| print(result) |
| else: |
| |
| interactive_mode(model, meta, args.device) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
