"""run.py — load the W4 checkpoint with the fused kernels and generate.
Usage: python run.py [your prompt...]        (GB10/CUDA box; first forward JIT-compiles ~45s)
"""
import pathlib
import sys
import time

HERE = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "kernels"))

from load_w4_checkpoint import load_fast

CKPT = HERE if (HERE / "w4_meta.json").exists() else HERE / "ckpt"
model, tok = load_fast(str(CKPT))                   # tier="bitexact" for token-identical output

prompt = " ".join(sys.argv[1:]) or "What is 17*23? Show your work."
ids = tok.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True,
                              return_tensors="pt", return_dict=True)["input_ids"].cuda()
t = time.time()
out = model.generate(input_ids=ids, max_new_tokens=512)
seq = out.sequences[0] if hasattr(out, "sequences") else out[0]
gen = seq[ids.shape[1]:]
print(tok.decode(gen, skip_special_tokens=True))
print(f"\n[{len(gen) / (time.time() - t):.0f} tok/s — includes ~45s JIT on the first run]")