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

This is the SFT (instruction-tuned) checkpoint -- unlike the base model, it
follows a chat format. This example does a single user turn -> assistant
response (no multi-turn history) to keep things simple.

Setup:
    pip install torch tiktoken rustbpe filelock kernels

Run from the root of this downloaded repo (where model_000849.pt lives):
    python inference.py --prompt "What is the capital of France?"
"""
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="Chat with lfm2-quantum-128m-sft")
parser.add_argument("--prompt", type=str, default="What is the capital of France?")
parser.add_argument("--max-tokens", type=int, default=256)
parser.add_argument("--temperature", type=float, default=0.8, help="0 = greedy decoding")
parser.add_argument("--top-k", type=int, default=10)
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_000849.pt + meta_000849.json live at the root of this repo.
model, _, meta = build_model(checkpoint_dir=".", step=849, device=device, phase="eval")
tokenizer = RustBPETokenizer.from_directory("tokenizer")

bos = tokenizer.get_bos_token_id()
user_start, user_end = tokenizer.encode_special("<|user_start|>"), tokenizer.encode_special("<|user_end|>")
assistant_start, assistant_end = tokenizer.encode_special("<|assistant_start|>"), tokenizer.encode_special("<|assistant_end|>")

conversation = [bos, user_start, *tokenizer.encode(args.prompt), user_end, assistant_start]

print(f"User: {args.prompt}\nAssistant: ", end="", flush=True)
for token_id in model.generate(
    conversation,
    max_tokens=args.max_tokens,
    temperature=args.temperature,
    top_k=args.top_k,
    seed=args.seed,
):
    if token_id == assistant_end:
        break
    print(tokenizer.decode([token_id]), end="", flush=True)
print()