Temperature and Top-K Sampling

Language models pick the next token from a probability distribution. Two settings shape it:

• Top-k: keep only the k tokens with the highest logits (if logits tie, the lower index ranks higher).
• Temperature T: divide the kept logits by T before softmax. Low T makes the choice more predictable; high T makes it more random.

To make sampling testable, the random number is given:

1. Keep the top k tokens, then list them by their original index (ascending).
2. Compute softmax(logit / T) over the kept tokens (subtract the largest before taking e^x).
3. Walk through the list adding probabilities; return the index of the first token where the running total becomes greater than random_value. If none does (rounding), return the last kept token.

With k = 1 this is greedy decoding: always the most likely token.

Implement sample_token(logits, temperature, k, random_value).

Input

logits: One score per token in the vocabulary.

temperature: T, greater than 0.

k: How many top tokens to keep.

random_value: A random number in [0, 1).

Output

Return the index of the chosen token.

Restricciones

1 ≤ k ≤ number of logits ≤ 1,000

0 < temperature ≤ 100, 0 ≤ random_value < 1

-100 ≤ each logit ≤ 100

Ejemplos

Entrada:
k = 1
logits = [1.0, 3.0, 2.0, 0.5]
temperature = 1.0
random_value = 0.99

Salida esperada:
expected = 1

Entrada:
long long k = 1;
std::vector<double> logits = {1.0, 3.0, 2.0, 0.5};
double temperature = 1.0;
double random_value = 0.99;

Salida esperada:
long long expected = 1;

Entrada:
long k = 1;
double[] logits = new double[] {1.0, 3.0, 2.0, 0.5};
double temperature = 1.0;
double random_value = 0.99;

Salida esperada:
long expected = 1;

Entrada:
k = 2
logits = [2.0, 1.0, 0.0]
temperature = 1.0
random_value = 0.8

Salida esperada:
expected = 1

Entrada:
long long k = 2;
std::vector<double> logits = {2.0, 1.0, 0.0};
double temperature = 1.0;
double random_value = 0.8;

Salida esperada:
long long expected = 1;

Entrada:
long k = 2;
double[] logits = new double[] {2.0, 1.0, 0.0};
double temperature = 1.0;
double random_value = 0.8;

Salida esperada:
long expected = 1;

Referencias

Este reto está adaptado para CodeUtility. Sigue el enunciado y las restricciones indicadas arriba.

• Hugging Face, How to generate text: decoding methods
• Goodfellow, Bengio and Courville, Deep Learning