#!/usr/bin/env python3
"""Szybki test inferencji GoLLeM v4 250M (seria migawek trajektorii).
    python generate.py "Polska to kraj" --n 60 --temp 0.8 --topk 40
    python generate.py "Polska to kraj" --ckpt ckpt_6000.safetensors   # konkretny etap
Domyslnie laduje NAJNOWSZA migawke. Wczesniejsze (2000/6000) sa slabsze - to migawki treningu.
"""
import argparse
from modeling_gollem import GollemGPT

ap = argparse.ArgumentParser()
ap.add_argument("prompt", nargs="?", default="Polska to kraj")
ap.add_argument("--path", default="./")
ap.add_argument("--ckpt", default=None, help="ckpt_<step>.safetensors; domyslnie najnowszy")
ap.add_argument("--n", type=int, default=60)
ap.add_argument("--temp", type=float, default=0.8)
ap.add_argument("--topk", type=int, default=40)
ap.add_argument("--device", default="cpu")
ap.add_argument("--seed", type=int, default=None)
a = ap.parse_args()

m = GollemGPT.from_pretrained(a.path, ckpt=a.ckpt, device=a.device)
print(f"[GoLLeM v4 250M base | {m._ckpt} | {sum(p.numel() for p in m.parameters())/1e6:.0f}M params]")
print(m.generate(a.prompt, max_new_tokens=a.n, temperature=a.temp, top_k=a.topk, seed=a.seed))