Easy Primitives

Greedy Decoding

Generate tokens one at a time by selecting the largest model score at each step. The supplied callback provides the scores; your function manages the sequence and the stopping rule.

Your task

Implement greedy_decode(logits_fn, prompt, max_tokens, eos_id).

prompt is a one-dimensional tensor of integer-valued token IDs.
logits_fn accepts the current sequence as a Python list of integers and returns a nonempty vector of finite logits, one score per vocabulary token. A token's ID is its position in that vector.
max_tokens is a nonnegative integer limiting newly generated tokens. Prompt tokens do not count toward this limit.
eos_id is the end-of-sequence ID. It may be outside the callback's vocabulary, in which case it cannot be generated.

Return a one-dimensional float32 tensor containing the prompt followed by all generated IDs. The float32 representation matches this exercise's runtime convention; the IDs passed to logits_fn must still be Python integers. Leave the prompt unchanged.

Generation loop

Start with an integer list copied from prompt. On each step:

Call logits_fn with the full current sequence.
Choose the index of the largest logit. If scores tie, choose the lowest index.
Append that ID to the sequence.
If the appended ID equals eos_id, stop. Keep that EOS token in the result.

For prompt [5], a callback that always favors token 2, max_tokens=3 and eos_id=99, return [5, 2, 2, 2].

If the callback favors token 7 while the sequence has fewer than three tokens, then favors EOS token 0, the same prompt grows as follows:

[5] -> [5, 7] -> [5, 7, 7] -> [5, 7, 7, 0]

With max_tokens=0, return the prompt without calling logits_fn. The stopping check applies to newly generated tokens; an EOS ID already in the prompt does not itself stop this loop.

What greedy selection does

Use torch.argmax or jnp.argmax. There is no random draw, and softmax is not needed to identify the largest finite logit. Fixed callback scores and the tie rule give a deterministic choice; the callback itself is responsible for any model randomness.

Greedy decoding chooses the best next token by score. It does not search for the highest-probability whole sequence, and it can repeat tokens. This exercise does not implement the language model, caching or beam search.

The diagram shows scores and their softmax probabilities. Changing positive temperature changes those probabilities but preserves the largest score's position, so greedy selection remains the same.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Choose the next token using the prefix so far

An autoregressive model predicts a distribution for the next token given the existing sequence. Greedy decoding picks the most likely token and repeats. Beam search keeps several candidate prefixes, expands each with possible next tokens, and retains the best candidates under its scoring rule.

Beam search considers more alternatives but can still discard a prefix that would later lead to the best complete sequence. Keep each candidate's state, score, and stopping status separate. A real model's next-token distribution depends on its prefix; a fixed table of scores by time step is only a simplified exercise.

Sequence log probability is a sum of token log probabilities. Since these terms are nonpositive, shorter sequences can be favored. Length normalization changes candidate rankings, so use the exact rule in the task.

Sample from a restricted distribution

Top-k sampling retains the k highest-scoring tokens. Nucleus, or top-p, sampling sorts tokens by probability and keeps the smallest prefix reaching the requested cumulative probability. Include the token that crosses the threshold.

For probabilities [0.6, 0.3, 0.1] and threshold 0.8, nucleus sampling retains the first two tokens because their total is 0.9. Renormalizing gives retained probabilities [2/3, 1/3]. Map any sampled sorted index back to its original vocabulary ID.

Temperature divides logits before softmax. A positive temperature changes how concentrated the distribution is. Check the task's end-of-sequence handling, maximum length, whether returned tokens include the prompt, and random-key usage.

Evaluate overlap without confusing it with meaning

BLEU-1 measures clipped unigram overlap with a reference and applies a brevity penalty. Repeating a reference word cannot earn unlimited credit. This unigram score ignores word order and does not establish semantic correctness. Follow the specified tokenization and empty-output conventions when computing it.

Detecting runtime…