A tiny language model
A language model predicts the next token from the ones before it. The smallest version uses characters and looks one character back: a bigram model. It is a table of scores, one row per current character, which training turns into next-character probabilities.
import torch
from torch import nn
vocab = 5
current = torch.tensor([0, 1, 2])
following = torch.tensor([1, 2, 3])
# row c: scores for what follows c
model = nn.Embedding(vocab, vocab)
logits = model(current)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, following)
print(logits.shape, loss.item())Generation repeats one step: look up the scores for the last character, turn them into probabilities, pick the next character, append it. Picking the most likely character every time gives the same text on each run; sampling with torch.multinomial gives variety.
Everything here, from encoding text to sampling, is what large models do too; they just look much further back, with attention.
Your task
Write encode(text) returning (ids, chars): chars is the sorted list of distinct characters, ids a long tensor of each character's index. Write train_bigram(ids, vocab, steps=300) returning an nn.Embedding(vocab, vocab) trained with Adam (lr=0.1) to predict each next id. Write generate(model, chars, start, length) that greedily appends the most likely next character length times and returns the whole string.
The explanation opens here once every check passes.