"""Lesson 12: quantize the tiny GPT to int8 (dynamic quantization) and compare size, speed and output."""
import io, time, torch
from tinygpt import TinyGPT, CharTokenizer, 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()
q = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

def size_kb(m):
    b = io.BytesIO(); torch.save(m.state_dict(), b); return b.tell() / 1024

print(f"fp32 {size_kb(model):.0f} KB  ->  int8 {size_kb(q):.0f} KB")
start = torch.tensor([tok.encode("The ")])
for name, m in [("fp32", model), ("int8", q)]:
    torch.manual_seed(3); t0 = time.time(); out = generate(m, start, 60, temperature=0.7)
    print(f"{name}: {time.time() - t0:.2f}s  " + repr(tok.decode(out[0].tolist())))