inference.py

"""
Minimal standalone inference example for lfm2-quantum-128m.

This is a BASE (pretrained, not instruction-tuned) checkpoint, so it does raw
text continuation, not chat.

Setup:
    pip install torch tiktoken rustbpe filelock kernels

Run from the root of this downloaded repo (where model_002162.pt lives):
    python inference.py --prompt "The history of quantum computing"
"""
import argparse
import torch

from nanochat.checkpoint_manager import build_model
from nanochat.tokenizer import RustBPETokenizer
from nanochat.common import autodetect_device_type

parser = argparse.ArgumentParser(description="Generate text from lfm2-quantum-128m")
parser.add_argument("--prompt", type=str, default="The meaning of life is")
parser.add_argument("--max-tokens", type=int, default=200)
parser.add_argument("--temperature", type=float, default=0.8, help="0 = greedy decoding")
parser.add_argument("--top-k", type=int, default=50)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--device-type", type=str, default="", choices=["cuda", "cpu", "mps"], help="empty = autodetect")
args = parser.parse_args()

device_type = args.device_type or autodetect_device_type()
device = torch.device(device_type)

# model_002162.pt + meta_002162.json live at the root of this repo.
model, _, meta = build_model(checkpoint_dir=".", step=2162, device=device, phase="eval")
tokenizer = RustBPETokenizer.from_directory("tokenizer")

bos = tokenizer.get_bos_token_id()
prompt_ids = [bos] + tokenizer.encode(args.prompt)

print(args.prompt, end="", flush=True)
for token_id in model.generate(
    prompt_ids,
    max_tokens=args.max_tokens,
    temperature=args.temperature,
    top_k=args.top_k,
    seed=args.seed,
):
    print(tokenizer.decode([token_id]), end="", flush=True)
print()