"""Lesson 8: sampling - temperature, top-k and top-p on the model from lesson 7 (run lesson 7 first)."""
import torch
from tinygpt import TinyGPT, generate, CharTokenizer
ck = torch.load("tinygpt.pt"); tok = CharTokenizer("".join(ck["chars"]))
model = TinyGPT(len(ck["chars"]), **ck["cfg"]); model.load_state_dict(ck["model"])
start = torch.tensor([tok.encode("The ")])
for name, kw in [("t=0.1", dict(temperature=0.1)), ("t=0.8", dict(temperature=0.8)), ("t=1.5", dict(temperature=1.5)),
                 ("top-k=5", dict(top_k=5)), ("top-p=0.9", dict(top_p=0.9))]:
    torch.manual_seed(7)
    print(name.rjust(10) + ": " + repr(tok.decode(generate(model, start, 60, **kw)[0].tolist())))