File size: 527 Bytes
c6c8db8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | """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())]))
|