| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
|
|
| |
| |
| |
| |
| |
| import torch |
| import torch.nn.functional as F |
|
|
| def get_candidate_tokens_logits(logits, beta): |
| """ |
| 直接基于logits筛选满足条件的token索引(不做归一化),返回格式同torch.topk的topk_idx |
| |
| 参数: |
| logits: 模型输出的未归一化分数,形状为 (batch_size, vocab_size) |
| beta: 阈值系数,范围通常为 [0, 1] |
| |
| 返回: |
| list[torch.Tensor]: 每个元素为对应批次样本的候选token索引张量, |
| 形状为 (num_candidates,)(num_candidates随样本变化) |
| """ |
| |
| max_logits, _ = logits.max(dim=-1, keepdim=True) |
| |
| |
| thresholds = beta * max_logits |
| |
| |
| masks = logits >= thresholds |
| |
| |
| candidate_indices = [] |
| for i in range(logits.size(0)): |
| indices = torch.where(masks[i])[0] |
| candidate_indices.append(indices) |
| |
| return torch.stack(candidate_indices, dim=0) |
|
|
| |
| logits = torch.tensor([ |
| [1.2, 3.5, 2.1, 0.8, 4.0], |
| [1.2, 3.5, 2.1, 0.8, 4.0], |
| ]) |
| beta = 0.6 |
|
|
| |
| candidates = get_candidate_tokens_logits(logits, beta) |
|
|
| print("样本1候选索引:", candidates) |
| |
| topk_vals = torch.gather(logits, dim=-1, index=candidates) |
| print(topk_vals) |
|
|
| def logits_entropy(logits: torch.Tensor, dim: int = -1, k: int = 8) -> torch.Tensor: |
| """ |
| 计算 logits 对应的熵 (Shannon entropy)。 |
| |
| 参数: |
| logits: torch.Tensor, shape [..., vocab_size] |
| dim: 计算概率分布的维度,一般是最后一维 (vocab_size) |
| 返回: |
| 熵值张量, shape [...] |
| """ |
| |
| probs = F.softmax(logits, dim=dim) |
| |
| |
| topk_probs = probs |
| log_probs = torch.log(topk_probs + 1e-12) |
| |
| entropy = -(topk_probs * log_probs).sum(dim=dim).item() |
| return entropy |
|
|
| print(logits_entropy(torch.tensor([[25.]]))) |