File size: 1,788 Bytes
68a7bb4 | 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 | """
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()
|