# 配套文章：03-training-and-generation.md
# 本轮新增示例未执行；环境和运行方式见上级 GUIDE.md。

import math

def probabilities(logits, temperature):
    # 温度必须为正；贪心选择应另外使用 argmax，而不是把温度设为零做除法。
    if temperature <= 0:
        raise ValueError("温度必须大于零")
    scaled = [x / temperature for x in logits]
    # 同时减去最大值保持分布不变，并减小指数溢出的风险。
    peak = max(scaled)
    masses = [math.exp(x - peak) for x in scaled]
    return [x / sum(masses) for x in masses]

def nucleus(probs, threshold):
    if not 0 < threshold <= 1:
        raise ValueError("top-p 必须位于 (0,1]")
    # 保留跨过阈值的那一项；至少有一个候选，不对单项概率设 threshold 门槛。
    chosen, total = [], 0.0
    for index in sorted(range(len(probs)), key=lambda i: probs[i], reverse=True):
        chosen.append(index)
        total += probs[index]
        if total >= threshold:
            break
    return {index: probs[index] / total for index in chosen}

logits = [2.0, 1.0, 0.0]
for temperature in [0.5, 1.0, 2.0]:
    print("温度：", temperature, "概率：", probabilities(logits, temperature))
print("top-p=0.8 后：", nucleus(probabilities(logits, 1.0), 0.8))