Next-token prediction: the only thing a model does

Given the tokens so far, a model gives a score to every token in its vocabulary for what comes next. Everything a chatbot does is built by repeating that one step.

import torch

inputs = tokenizer("The parcel", return_tensors="pt")
with torch.no_grad():
    logits = model(**inputs).logits

print(logits.shape)

return_tensors="pt" returns the ids as a PyTorch tensor, a grid of numbers, which is what the model takes. torch.no_grad() says you are only asking for an answer, not training, which saves memory.

The shape reads as: 1 input, 2 token positions, one score for each vocabulary entry at each position. The scores are called logits. Only the last position matters for what comes next.

Scores into probabilities

probs = torch.softmax(logits[0, -1], dim=-1)
top = torch.topk(probs, 5)

for p, i in zip(top.values, top.indices):
    print(f"{tokenizer.decode(i)!r:14} {p.item():.3f}")

softmax turns the scores into probabilities that add up to 1. topk keeps the five largest. !r:14 prints each token with quotes, so the leading space shows, padded to 14 characters.

No single continuation is certain. The model thinks ' delivery' is the best bet, at about one chance in eight, with many others close behind. The answer you finally read depends on how one of these is picked, which is the next two lessons.

A longer context

inputs = tokenizer("Sorry, your parcel", return_tensors="pt")
with torch.no_grad():
    probs = torch.softmax(model(**inputs).logits[0, -1], dim=-1)
top = torch.topk(probs, 5)
print([(tokenizer.decode(i), round(p.item(), 3)) for p, i in zip(top.values, top.indices)])

Two more words in front change the whole list. The model reads everything before the next token, and every token of the prompt shifts these probabilities. That is the entire mechanism a prompt works through.