Amostragem com temperatura e top-k

Modelos de linguagem escolhem o próximo token a partir de uma distribuição de probabilidade. Dois ajustes a moldam:

• 'Top-k: mantenha apenas os k tokens com os maiores logits (com logits iguais, o menor índice fica à frente).'
• 'Temperatura T: divida os logits mantidos por T antes do softmax. Um T baixo torna a escolha mais previsível; um T alto, mais aleatória.'

Para que a amostragem possa ser testada, o número aleatório é fornecido:

1. Mantenha os tokens do top-k e liste-os pelo índice original (crescente).
2. Calcule softmax(logit / T) sobre os tokens mantidos (subtraia o maior antes de calcular e^x).
3. Percorra a lista somando as probabilidades; retorne o índice do primeiro token em que o total acumulado ultrapassa random_value. Se isso nunca acontecer (arredondamento), retorne o último token mantido.

Com k = 1 é a decodificação gulosa: sempre o token mais provável.

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

Entrada

logits: Uma pontuação por token do vocabulário.

temperature: T, maior que 0.

k: Quantos dos melhores tokens manter.

random_value: Um número aleatório em [0, 1).

Saída

Retorne o índice do token escolhido.

Restrições

1 ≤ k ≤ número de logits ≤ 1.000

0 < temperature ≤ 100, 0 ≤ random_value < 1

-100 ≤ cada logit ≤ 100

Exemplos

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

Saída 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;

Saída 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;

Saída esperada:
long expected = 1;

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

Saída 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;

Saída 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;

Saída esperada:
long expected = 1;

Referências

Este desafio foi adaptado para o CodeUtility. Siga o enunciado e as restrições acima.

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