"""Lesson 9: evaluate - average loss and perplexity on held-out text, plus fixed test prompts."""
import math, torch
from tinygpt import TinyGPT, CharTokenizer, get_batch, generate
ck = torch.load("tinygpt.pt"); tok = CharTokenizer("".join(ck["chars"]))
model = TinyGPT(len(ck["chars"]), **ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
text = open("sample.txt", encoding="utf-8").read(); data = torch.tensor(tok.encode(text)); val = data[int(0.9 * len(data)):]
torch.manual_seed(0)
with torch.no_grad():
    losses = [model(*get_batch(val, 64, 32))[1].item() for _ in range(20)]
avg = sum(losses) / len(losses)
print(f"validation loss {avg:.3f}  perplexity {math.exp(avg):.2f}  (random guessing = {tok.vocab_size})")
for p in ["Excel ", "SQL ", "GST "]:
    torch.manual_seed(1)
    print(repr(tok.decode(generate(model, torch.tensor([tok.encode(p)]), 50, temperature=0.7)[0].tolist())))
