Scale's GenAI Data Engine may need to sample the next token from a language model's logits. Implement temperature scaling and nucleus, or top-p, sampling from scratch without using machine learning libraries.
Given a list of token logits, a positive temperature, a cumulative probability threshold top_p, and a random seed, return one sampled token index.
Implement sample_next_token(logits, temperature, top_p, seed), where logits is a non-empty list of real numbers. First divide every logit by temperature, then compute a numerically stable softmax distribution. Sort tokens by descending probability, breaking ties by ascending token index. Keep the smallest prefix whose cumulative probability is at least top_p, renormalize the retained probabilities, and draw one token using random.Random(seed).
Return the selected token's original index as an integer. Do not mutate logits.
def sample_next_token(logits, temperature, top_p, seed):