Nucleus Sampling Generation Loop

Build a token sequence by applying temperature scaling and nucleus filtering to a callback's scores, then drawing one token at each step. The callback supplies the model behavior; you implement the generation loop.

Your task

Implement generate(logits_fn, prompt, max_tokens, eos_id, temperature, top_p, seed). This exercise is graded in PyTorch.

• logits_fn accepts the full current sequence as a Python list of integers. It returns a nonempty float32 CPU vector of finite logits, indexed by token ID. The runtime supplies the callable itself, not a name you must resolve.
• prompt is a one-dimensional tensor of integer-valued starting IDs.
• max_tokens is a nonnegative limit on new tokens, excluding the prompt.
• eos_id is the ID that ends generation when newly sampled.
• temperature is a finite positive number; top_p satisfies 0 < top_p <= 1.
• seed is an integer. Step numbering starts at zero.

Return a one-dimensional int64 tensor containing the prompt followed by all new IDs. Leave the prompt unchanged. An empty prompt and zero new tokens produce an empty int64 vector.

One generation step

Call logits_fn with the current integer list. Divide its logits by temperature, then compute softmax by subtracting the largest scaled score before exponentiating. A lower positive temperature concentrates probability on larger scores; a higher one makes probabilities more uniform. It does not by itself guarantee diverse or coherent output.

For top_p below 1, sort probabilities in descending order, breaking ties by lower token ID. Retain the shortest prefix whose cumulative mass reaches or exceeds top_p. Include the threshold-crossing token and always retain the first. For top_p=1, keep the full distribution without filtering.

Restore the retained probabilities to their original vocabulary positions, set excluded entries to zero and divide by the retained total. Preserve the full vocabulary-sized vector for sampling.

Use torch.multinomial(filtered, num_samples=1, generator=gen) with a fresh gen = torch.Generator().manual_seed(seed + step) on each step. Reusing one generator across steps or sampling a reordered vector can produce different seeded outputs. Return the original vocabulary ID, not its sorted position.

Append the sampled ID, then stop if it equals eos_id. Keep that EOS token. Otherwise continue until max_tokens new tokens have been added. Do not stop merely because the original prompt contains EOS. With max_tokens=0, return the prompt without calling the callback or sampling.

For prompt [5], a callback that favors token 7 until the sequence has three tokens and then favors EOS token 0 can produce [5, 7, 7, 0]. The final 0 is included. The supplied peaked fixtures make these particular draws predictable; ordinary distributions can produce several valid-looking sequences, so the specified seed recipe matters for matching the tests.

The diagram illustrates temperature's effect on probabilities. The loop then applies nucleus filtering and makes a seeded draw from those probabilities.