解码采样器

题目描述

大语言模型生成文本时，每一步都要从词表上的概率分布中选出一个 token。真实推理框架
（如 HuggingFace generate、vLLM）支持若干采样参数来调节输出的随机性与质量：
温度（temperature）、top-k、top-p（nucleus）与重复惩罚（repetition penalty）。

本题要求你实现这条完整的采样流水线：给定一组 logits 与采样参数，输出过滤后的
概率分布以及本次采样选中的 token。

为消除随机性，采样不使用随机数发生器，而是由输入给出一个均匀随机数 $u\in [0,1)$，
你用逆变换采样（inverse transform sampling）确定结果。因此对相同输入，输出必须完全一致。

输入格式

一行 JSON 对象（input_str 即该 JSON 文本），字段如下：

字段	类型	说明
logits	数组	长度 $V\ge 1$，第 $i$ 项是 token $i$ 的未归一化对数概率。元素为数字；字符串 "-inf" 表示负无穷
temperature	数字	$\ge 0$。为 $0$ 时表示贪心解码
top_k	整数	$\le 0$ 表示不做 top-k 过滤；否则保留概率最高的 $k$ 个
top_p	数字	$\in (0,1]$，nucleus 采样的累计概率阈值
repetition_penalty	数字	$\ge 1.0$，重复惩罚系数
previous_tokens	整数数组	已出现过的 token 索引，可能为空、可能含重复项、可能含越界值
u	数字	$\in [0,1)$，用于逆变换采样的均匀随机数

输出格式

一行 JSON 对象：

{"token_id": 2, "probs": [0.09, 0.24, 0.67]}

字段	说明
token_id	本次选中的 token 索引；无可用 token 时为 -1
probs	长度必须等于 $V$ 的概率分布；被过滤掉的 token 概率为 0.0

允许在结果行之前打印任意调试信息（评测时只取最后一行非空内容作为结果）。

规则（严格按此顺序）

记输入 logits 为 <math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><msub><mi mathvariant="normal">ℓ</mi><mn>0</mn></msub><mo separator="true">,</mo><mo>…</mo><mo separator="true">,</mo><msub><mi mathvariant="normal">ℓ</mi><mrow><mi>V</mi><mo>−</mo><mn>1</mn></mrow></msub></mrow><annotation encoding="application/x-tex">\ell_0, \dots, \ell_{V-1}</annotation></semantics></math>ℓ0​,…,ℓV−1​。

第 1 步：重复惩罚
令 $P$ = previous_tokens 中所有合法（$0\le t<V$）且去重后的索引集合。
对每个 $t\in P$：

<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>t</mi></msub><mo>←</mo><mrow><mo fence="true">{</mo><mtable rowspacing="0.36em" columnalign="left left" columnspacing="1em"><mtr><mtd><mstyle scriptlevel="0" displaystyle="false"><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>t</mi></msub><mi mathvariant="normal">/</mi><mtext>penalty</mtext><mo separator="true">,</mo></mrow></mstyle></mtd><mtd><mstyle scriptlevel="0" displaystyle="false"><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>t</mi></msub><mo>></mo><mn>0</mn></mrow></mstyle></mtd></mtr><mtr><mtd><mstyle scriptlevel="0" displaystyle="false"><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>t</mi></msub><mo>×</mo><mtext>penalty</mtext><mo separator="true">,</mo></mrow></mstyle></mtd><mtd><mstyle scriptlevel="0" displaystyle="false"><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>t</mi></msub><mo>≤</mo><mn>0</mn></mrow></mstyle></mtd></mtr></mtable></mrow></mrow><annotation encoding="application/x-tex">\ell_t \leftarrow
\begin{cases}
\ell_t / \text{penalty}, & \ell_t > 0 \
\ell_t \times \text{penalty}, & \ell_t \le 0
\end{cases}</annotation></semantics></math>ℓt​←{ℓt​/penalty,ℓt​×penalty,​ℓt​>0ℓt​≤0​

若 repetition_penalty == 1.0，此步不改变任何值。

第 2 步：温度

• 若 temperature == 0：贪心解码。取调整后 logits 的最大值所在索引；
  若存在并列最大值，取最小索引。输出该 token 的 one-hot 分布
  （选中项概率 1.0，其余 0.0）。此时忽略 top_k、top_p 与 u。
  若所有 logits 均为 $-\infty$，则输出 token_id = -1 且 probs 全为 0.0。
• 否则，令 <math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>i</mi></msub><mo>←</mo><msub><mi mathvariant="normal">ℓ</mi><mi>i</mi></msub><mi mathvariant="normal">/</mi><mtext>temperature</mtext></mrow><annotation encoding="application/x-tex">\ell_i \leftarrow \ell_i / \text{temperature}</annotation></semantics></math>ℓi​←ℓi​/temperature，继续第 3 步。

第 3 步：top-k 过滤
若 top_k > 0：保留 <math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><msub><mi mathvariant="normal">ℓ</mi><mi>i</mi></msub></mrow><annotation encoding="application/x-tex">\ell_i</annotation></semantics></math>ℓi​ 最大的 $\min (\text{top\_k},V)$ 个索引；
排序并列时优先较小索引。未保留的索引视为 $-\infty$。
若 top_k <= 0 或 top_k >= V，不裁剪。

第 4 步：top-p（nucleus）过滤
若 top_p < 1.0：

1. 把当前保留的索引按「logit 降序、并列时索引升序」排序；
2. 对该序列做 softmax 得到概率；
3. 从前往后累加，保留到累计概率 $\ge$ top_p 的那个 token 为止（含该 token）；
4. 至少保留 1 个 token。

若 top_p >= 1.0，不裁剪。

第 5 步：最终分布
对第 4 步后保留的索引集合做数值稳定的 softmax（减去集合内最大值后取指数），
得到长度 $V$ 的 probs，未保留位置为 0.0。

若所有保留位置的 logit 均为 $-\infty$（概率总和为 $0$），输出 token_id = -1 且 probs 全为 0.0。

第 6 步：逆变换采样
按 token 索引升序（$0,1,\dots ,V-1$）累加 probs，取第一个使累计和严格大于 $u$ 的 token 作为 token_id。

注意第 4 步的排序（按 logit 降序）与第 6 步的累加顺序（按索引升序）不同，
不要混用。

示例

示例 1

输入：

{"logits": [2.0, 1.0, 0.0, -1.0], "temperature": 1.0, "top_k": 2, "top_p": 1.0, "repetition_penalty": 1.0, "previous_tokens": [], "u": 0.3}

输出：

{"token_id": 0, "probs": [0.7310585786300049, 0.2689414213699951, 0.0, 0.0]}

说明：top-k=2 保留索引 0、1；softmax 后 $P(0)=0.7311$；$u=0.3<0.7311$，
故按索引升序累加时第一个超过 $u$ 的是 token 0。

示例 2

输入：

{"logits": [0.5, 3.0, 1.0], "temperature": 0.0, "top_k": 0, "top_p": 1.0, "repetition_penalty": 1.0, "previous_tokens": [], "u": 0.99}

输出：

{"token_id": 1, "probs": [0.0, 1.0, 0.0]}

说明：温度为 0 走贪心，最大值在索引 1，输出 one-hot 分布（u 被忽略）。

示例 3

输入：

{"logits": [1.0, 1.0, 1.0], "temperature": 1.0, "top_k": 0, "top_p": 1.0, "repetition_penalty": 2.0, "previous_tokens": [1], "u": 0.5}

输出：

{"token_id": 1, "probs": [0.38365173119055074, 0.23269653761889864, 0.38365173119055074]}

说明：token 1 的 logit 为正，被除以 2 变成 0.5，因此概率低于另外两个；
$u=0.5$ 时按索引升序累加：$0.3837<0.5$，$0.6163>0.5$，故选中 token 1。

数据范围与限制

• $1\le V\le 1000$
• temperature $\in [0,100]$（为 0 时按贪心处理）
• top_k <math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mo>∈</mo><mo stretchy="false">[</mo><mo>−</mo><mn>1</mn><mo separator="true">,</mo><msup><mn>10</mn><mn>6</mn></msup><mo stretchy="false">]</mo></mrow><annotation encoding="application/x-tex">\in [-1, 10^6]</annotation></semantics></math>∈[−1,106]
• top_p $\in (0,1]$
• repetition_penalty $\in [1.0,100.0]$
• logits 元素 <math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mo>∈</mo><mo stretchy="false">[</mo><mo>−</mo><msup><mn>10</mn><mn>6</mn></msup><mo separator="true">,</mo><msup><mn>10</mn><mn>6</mn></msup><mo stretchy="false">]</mo></mrow><annotation encoding="application/x-tex">\in [-10^6, 10^6]</annotation></semantics></math>∈[−106,106]，或字符串 "-inf"
• 时间限制：$2000\text{ms}$ / 用例；内存限制：$256\text{MB}$

单用例超时（记 TimeLimitExceeded）与选手代码抛异常（记 RuntimeError）都只影响
该用例（该用例 0 分），评测继续执行其余用例——一条用例崩溃不会连带丢掉其余用例的
得分；只有 Solution 进程崩溃、评测通道断开等评测侧故障才会使整份评测失败。

评分

• 正式得分只来自隐藏用例，各用例等权，满分 100 分。
• 判定标准：token_id 完全一致，且 probs 长度一致、逐元素绝对误差 <math xmlns="http://www.w3.org/1998/Math/MathML"><semantics><mrow><mo>≤</mo><msup><mn>10</mn><mrow><mo>−</mo><mn>9</mn></mrow></msup></mrow><annotation encoding="application/x-tex">\le 10^{-9}</annotation></semantics></math>≤10−9。
• 题面中的 3 个示例会作为可见用例运行并展示调试信息，但不计分。