Generating Text with a Simple Loop

Earlier in this module, the text generation loop produced each new token by looking only at the last time step of the model's output and choosing the single most likely token. That approach is simple and predictable, but it always makes the same choice when given the same input. A generator that always picks the top candidate tends to produce repetitive and narrow text. This article shows how a modified version of that loop, called generate, adds two controls to introduce variety while keeping the whole mechanism inside one function.

The modified function is the earlier generate_text_simple loop, extended with two controls. Temperature scaling changes how random the choice is. Top-k sampling restricts the set of tokens the choice can come from. Both controls live inside the same loop, so you can read the full mechanism in one place. The best way to start is to see what the modified function produces, then trace how each line produces it.

The first part of the loop body is the same as before. It crops the context to context_size, runs the model, and keeps only the logits for the last time step. Everything after that line is new, and that is where the added diversity comes from. Each pass through the loop follows one pipeline: top-k filtering runs first, then temperature scaling, then softmax converts the logits into probabilities, and finally torch.multinomial samples one token from those probabilities.

The order matters because each step changes the input to the next one. Top-k filtering decides which candidates remain eligible. Temperature scaling then reshapes the logits of those candidates before softmax turns them into probabilities. Sampling is the final step, and it draws from the probabilities that the earlier steps produced. Reading the code in this order, rather than line by line in isolation, makes the purpose of each control clear.

Temperature scaling works on the logits themselves. Let z stand for the logits of the candidate tokens and T stand for the temperature. The function passes the scaled logits z' to softmax, where the scaling is defined as follows.

In the code, this corresponds to the line logits = logits / temperature. Temperature is the divisor, and the scaled logits are what softmax turns into probabilities for sampling. The purpose of this division is to change how sharply the probabilities differ after softmax. The source does not show the probability values directly, so the exact effect depends on the logits at each step. A useful way to read it is that changing the divisor changes the gaps between the logits before they are converted into probabilities.

Top-k sampling restricts candidates to the k largest logits. The cutoff is the smallest of the top k logits. Any logit below that cutoff is replaced with negative infinity. Softmax is applied later in the function, and a token whose logit is negative infinity receives no probability, so sampling cannot select it. This narrows the candidate set to the top k tokens without removing the remaining tokens from the vocabulary arrays. Their positions still exist; they simply carry no probability.

In the diagram, the two largest logits, 3.0 and 2.0, survive. The cutoff is 2.0, the smaller of the two. Every other logit falls below that cutoff and is replaced with negative infinity. Those positions keep their place in the vocabulary, but they receive zero probability after softmax. The point of the toy example is the mechanism, not the specific numbers.

The temperature value decides which selection path runs. With temperature greater than 0.0, the function samples a token from the probabilities. With temperature 0.0 or lower, the function skips the scaling and sampling lines and calls torch.argmax, which selects the token with the highest logit. The earlier generator always made this greedy choice, so the greedy branch preserves that earlier behavior as a special case of the modified function.

After a token is selected by either path, the function checks whether it equals eos_id, the end-of-sequence token ID. When it does, the break statement exits the loop and the function returns the sequence built so far, so generation can stop before max_new_tokens is reached. If the end-of-sequence token never appears, the loop runs for max_new_tokens iterations, and each pass appends one token with torch.cat.

The example runs the full mechanism on a case small enough to check by hand. It also shows that top-k filtering is a restriction on candidates, while temperature decides how the remaining candidates are weighted before sampling.

The source runs the modified function with torch.manual_seed(123) set first. That call fixes the random number generator, so the sampling step behaves reproducibly on that run. The call also passes top_k=25 and temperature=1.4, so both controls are active. The output is the prompt followed by up to 15 newly generated tokens. The text reads as plausible words that drift in meaning, which is the variety the modified function is designed to allow.

Read the run as a chain of branch decisions. Each pass checks temperature, which is 1.4, so the sampling path is taken on every iteration. Each pass also checks whether the chosen token equals eos_id. If that happens, the loop stops early. Otherwise, the loop continues until max_new_tokens iterations have added 15 tokens. The seed governs the random draws inside torch.multinomial, so repeating the same seeded call on the same setup should reproduce the same sequence of choices.

Determinism depends on which branch runs. On the greedy path, temperature is 0.0 or lower, and the function calls torch.argmax. That selects the token with the highest logit, and given the same logits it always selects the same token. On the sampling path, temperature is greater than 0.0, and torch.multinomial draws a random token. Two runs with the same settings and no fixed seed can therefore diverge, while a run with a fixed seed reproduces its draws. Top-k filtering does not by itself introduce randomness; it only narrows which tokens the sampling step can draw from.