| """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())])) | |