"""Minimal loader. Requires: torch, tokenizers."""
import torch
from tokenizers import Tokenizer
from simloop.stack import StackConfig, StackedLoop

ck = torch.load("model.pt", map_location="cpu", weights_only=False)
model = StackedLoop(StackConfig(**ck["model_cfg"]))
model.load_state_dict(ck["model"])
model.eval()

tok = Tokenizer.from_file("tokenizer.json")
ids = torch.tensor([tok.encode("The capital of France is").ids])
with torch.no_grad():
    logits = model(ids, K=1)
print(tok.decode([int(logits[0, -1].argmax())]))