Temperature Sampling

Task

Implement temperature_sampling(step_logits: list[list[float]], random_values: list[float], temperature: float, eos_token_id: int, max_new_tokens: int) -> list[int]. Generate tokens autoregressively by applying temperature-scaled stable softmax to each logit row and using the corresponding random value for inverse-CDF sampling.

Requirements

• Each row in step_logits contains the vocabulary logits for one autoregressive generation step.
• random_values contains deterministic uniform draws in [0, 1), one for each available generation step.
• Subtract the largest row logit, divide each difference by the positive temperature, and then exponentiate.
• Normalize the exponentials into a probability distribution over the vocabulary.
• Select the first token whose cumulative probability is strictly greater than the step's random value.
• If floating-point rounding leaves no cumulative value greater than the draw, select the final token.
• Append the sampled token before checking whether it equals eos_token_id, then stop immediately on EOS.
• Generate at most max_new_tokens tokens and stop when either logits or random values are exhausted.
• Return only newly generated token IDs and do not mutate either input.
• Assume temperature is greater than zero, every consumed logit is finite, every consumed row is non-empty, and every random value lies in [0, 1).

Example

step_logits = [
    [0.0, 1.0986122887],
    [1.0986122887, 0.0],
    [0.0, 0.0],
]
random_values = [0.2, 0.8, 0.4]

temperature_sampling(
    step_logits,
    random_values,
    temperature=1.0,
    eos_token_id=9,
    max_new_tokens=3,
)
# [0, 1, 0]