generate.py

"""
generate.py — Text Generation Inference Script
================================================

Generate text from a trained GPT-2 checkpoint or a HuggingFace pre-trained model.

From checkpoint::

    python generate.py --ckpt checkpoints/best.pt --prompt "The brain is"

From HuggingFace::

    python generate.py --pretrained small --prompt "Neurons communicate via"

Interactive REPL mode::

    python generate.py --pretrained small --interactive
"""

from __future__ import annotations

import argparse
import sys
import logging

import torch
import tiktoken

from src.config.model_config import GPTConfig, ModelScale
from src.config.system_config import SystemConfig
from src.model.gpt import GPT
from src.utils import setup_logging


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="NanoGPT-Production — Text Generation")
    source = p.add_mutually_exclusive_group(required=True)
    source.add_argument("--ckpt",       type=str, help="Path to a local .pt checkpoint.")
    source.add_argument("--pretrained", type=str, choices=["small", "medium", "large", "xl"],
                        help="Load from HuggingFace OpenAI GPT-2 weights.")

    p.add_argument("--prompt",       type=str,   default="Hello, I am a language model,")
    p.add_argument("--max_tokens",   type=int,   default=200)
    p.add_argument("--temperature",  type=float, default=1.0)
    p.add_argument("--top_k",        type=int,   default=50)
    p.add_argument("--top_p",        type=float, default=0.95)
    p.add_argument("--num_samples",  type=int,   default=1)
    p.add_argument("--interactive",  action="store_true", help="Enter REPL mode.")
    p.add_argument("--device",       type=str,   default=None,
                   help="Override device (e.g. 'cpu', 'cuda:1').")
    return p.parse_args()


def load_model(args: argparse.Namespace, device: str) -> GPT:
    """Load GPT model from checkpoint or HuggingFace weights."""
    if args.ckpt:
        logging.getLogger(__name__).info(f"Loading checkpoint: {args.ckpt}")
        ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
        # Restore full architecture from saved config including Llama-3 flags
        cfg_dict = ckpt["config"]
        cfg = GPTConfig(**{
            k: cfg_dict[k]
            for k in GPTConfig.__dataclass_fields__
            if k in cfg_dict
        })
        model = GPT(cfg)
        model.load_state_dict(ckpt["model_state_dict"])
    else:
        scale_map = {
            "small":  ModelScale.SMALL,
            "medium": ModelScale.MEDIUM,
            "large":  ModelScale.LARGE,
            "xl":     ModelScale.XL,
        }
        model = GPT.from_pretrained(scale_map[args.pretrained])

    model.eval()
    model.to(device)
    return model


def generate_text(
    model: GPT,
    enc: "tiktoken.Encoding",
    prompt: str,
    max_tokens: int,
    temperature: float,
    top_k: int,
    top_p: float,
    device: str,
    num_samples: int = 1,
) -> list[str]:
    """Generate text completions for a given prompt.

    Args:
        model: Loaded GPT model in eval mode.
        enc: GPT-2 tiktoken encoder.
        prompt: Input text prompt.
        max_tokens: Number of new tokens to generate.
        temperature: Sampling temperature.
        top_k: Top-k filtering (0 = disabled).
        top_p: Nucleus sampling probability (1.0 = disabled).
        device: Target device.
        num_samples: Number of independent generations.

    Returns:
        List of generated text strings (including the prompt prefix).
    """
    tokens = enc.encode(prompt)
    x = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)
    x = x.repeat(num_samples, 1)  # (num_samples, T)

    with torch.no_grad():
        output = model.generate(
            x,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_k=top_k if top_k > 0 else None,
            top_p=top_p if top_p < 1.0 else None,
        )

    results = []
    for i in range(num_samples):
        generated = enc.decode(output[i].tolist())
        results.append(generated)
    return results


def repl(model: GPT, enc, args, device: str) -> None:
    """Interactive REPL for continuous text generation."""
    print("\n" + "=" * 60)
    print("NanoGPT Interactive Mode  |  Ctrl+C or type 'quit' to exit")
    print("=" * 60)

    while True:
        try:
            prompt = input("\nPrompt > ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\nExiting.")
            break

        if prompt.lower() in {"quit", "exit", "q"}:
            break
        if not prompt:
            continue

        texts = generate_text(
            model, enc, prompt,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
            top_p=args.top_p,
            device=device,
            num_samples=1,
        )
        print("\n" + "-" * 60)
        print(texts[0])
        print("-" * 60)


def main() -> None:
    args = parse_args()
    setup_logging(rank=0)
    logger = logging.getLogger(__name__)

    # Device selection
    if args.device:
        device = args.device
    elif torch.cuda.is_available():
        device = "cuda"
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"

    logger.info(f"Using device: {device}")

    model = load_model(args, device)
    enc = tiktoken.get_encoding("gpt2")

    if args.interactive:
        repl(model, enc, args, device)
        return

    # Single-shot generation
    print(f"\nPrompt: {args.prompt!r}")
    print(f"Generating {args.max_tokens} tokens × {args.num_samples} sample(s)...")
    print("=" * 60)

    texts = generate_text(
        model, enc, args.prompt,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        device=device,
        num_samples=args.num_samples,
    )

    for i, text in enumerate(texts, 1):
        print(f"\n[Sample {i}]\n{text}")
        print("-" * 60)


if __name__ == "__main__":
    main()