LLM模型贪婪、温度、Top-k、核采样方式的区别（附代码与示例）

class GreedySampler(Sampler):
    def __call__(self, logits: torch.Tensor):
        return logits.argmax(dim=-1)

class TemperatureSampler(Sampler):
    def __init__(self, temperature: float = 1.0):
        self.temperature = temperature

    def __call__(self, logits: torch.Tensor):
        dist = Categorical(logits=logits / self.temperature)
        return dist.sample()

class TopKSampler(Sampler):
    def __init__(self, k: int, sampler: Sampler):
        self.k = k
        self.sampler = sampler

    def __call__(self, logits: torch.Tensor):
        zeros = logits.new_ones(logits.shape) * float('-inf')
        values, indices = torch.topk(logits, self.k, dim=-1)
        zeros.scatter_(-1, indices, values)
        return self.sampler(zeros)

class NucleusSampler(Sampler):
    """
    ## Nucleus 采样器

    Nucleus 采样器根据给定的概率 p 选择词汇的一个子集，并从中进行采样。
    """

    def __init__(self, p: float, sampler: Sampler):
        """
        ### 初始化

        :param p: 要选择的令牌概率之和，即 p 值。
        :param sampler: 用于从选定令牌中进行采样的采样器。
        """
        # 保存 p 值
        self.p = p
        # 保存采样器
        self.sampler = sampler
        # 初始化 softmax 层，用于将 logits 转换为概率
        self.softmax = nn.Softmax(dim=-1)

    def __call__(self, logits: torch.Tensor):
        """
        ### 从 logits 中进行 Nucleus 采样

        :param logits: 输入的 logits 张量，形状为 (batch_size, num_tokens)。
        :return: 采样得到的令牌索引，形状为 (batch_size,)。
        """
        # 获取概率 P(x_i | x_1:i-1)
        probs = self.softmax(logits)

        # 按降序对概率进行排序，并获取排序后的索引
        sorted_probs, indices = torch.sort(probs, dim=-1, descending=True)

        # 按排序顺序获取概率的累积总和
        cum_sum_probs = torch.cumsum(sorted_probs, dim=-1)

        # 找出累积总和小于 p 的令牌
        nucleus = cum_sum_probs < self.p

        # 在前面加一个 True，这样我们可以在累积概率小于 p 的最小令牌数量之后添加一个令牌
        nucleus = torch.cat([nucleus.new_ones(nucleus.shape[:-1] + (1,)), nucleus[..., :-1]], dim=-1)

        # 获取对数概率并掩盖非核部分
        sorted_log_probs = torch.log(sorted_probs)
        sorted_log_probs[~nucleus] = float('-inf')

        # 使用采样器从排序后的对数概率中进行采样
        sampled_sorted_indexes = self.sampler(sorted_log_probs)

        # 获取实际的索引
        res = indices.gather(-1, sampled_sorted_indexes.unsqueeze(-1))

        # 返回采样得到的令牌索引
        return res.squeeze(-1)