Сэмплирование с температурой и top-k

Языковые модели выбирают следующий токен из распределения вероятностей. Его формируют две настройки:

• 'Top-k: оставьте только k токенов с наибольшими логитами (при равных логитах выше стоит меньший индекс).'
• 'Температура T: разделите оставленные логиты на T перед softmax. Низкая T делает выбор более предсказуемым, высокая — более случайным.'

Чтобы сэмплирование можно было проверить, случайное число передаётся на вход:

1. Оставьте top-k токенов и перечислите их по исходному индексу (по возрастанию).
2. Вычислите softmax(logit / T) по оставленным токенам (вычтите наибольшее значение перед вычислением e^x).
3. Пройдите по списку, суммируя вероятности; верните индекс первого токена, на котором накопленная сумма станет больше random_value. Если этого не произошло (из-за округления), верните последний оставленный токен.

При k = 1 это жадное декодирование: всегда самый вероятный токен.

Реализуйте sample_token(logits, temperature, k, random_value).

Входные данные

logits: По одной оценке на каждый токен словаря.

temperature: T, больше 0.

k: Сколько лучших токенов оставить.

random_value: Случайное число в [0, 1).

Выходные данные

Верните индекс выбранного токена.

Ограничения

1 ≤ k ≤ число логитов ≤ 1 000

0 < temperature ≤ 100, 0 ≤ random_value < 1

-100 ≤ каждый логит ≤ 100

Примеры

Ввод:
k = 1
logits = [1.0, 3.0, 2.0, 0.5]
temperature = 1.0
random_value = 0.99

Ожидаемый вывод:
expected = 1

Ввод:
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;

Ожидаемый вывод:
long long expected = 1;

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

Ожидаемый вывод:
long expected = 1;

Ввод:
k = 2
logits = [2.0, 1.0, 0.0]
temperature = 1.0
random_value = 0.8

Ожидаемый вывод:
expected = 1

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

Ожидаемый вывод:
long long expected = 1;

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

Ожидаемый вывод:
long expected = 1;

Ссылки

Это задание адаптировано для CodeUtility. Следуйте условию и ограничениям выше.

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