(logits: torch.Tensor, metrics: Dict[str, torch.Tensor],
gen_tokens: torch.Tensor, n_samples: int,
base_temp: float = 0.666, base_top_p: float = 0.90, base_top_k: int = 40, base_min_p: float = 0.03,
generator: torch.Generator = None)
| 79 | return next_token_g.to(torch.int32) |
| 80 | |
| 81 | def adaptive_sample(logits: torch.Tensor, metrics: Dict[str, torch.Tensor], |
| 82 | gen_tokens: torch.Tensor, n_samples: int, |
| 83 | base_temp: float = 0.666, base_top_p: float = 0.90, base_top_k: int = 40, base_min_p: float = 0.03, |
| 84 | generator: torch.Generator = None) -> torch.Tensor: |
| 85 | logits_uncertainty = metrics["logits_entropy"] + metrics["logits_varentropy"] |
| 86 | attn_uncertainty = metrics["attn_entropy"] + metrics["attn_varentropy"] |
| 87 | |
| 88 | temperature = base_temp * (1 + 0.3 * logits_uncertainty + 0.2 * attn_uncertainty - 0.2 * metrics["agreement"]) |
| 89 | top_p = torch.clamp(base_top_p * (1 + 0.1 * metrics["attn_varentropy"]), 0.1, 1.0) |
| 90 | top_k = int(torch.clamp( |
| 91 | torch.round(torch.tensor(base_top_k) * (1 + 0.3 * metrics["interaction_strength"].item() - 0.2 * metrics["agreement"].item())), |
| 92 | min=1, |
| 93 | max=100 |
| 94 | ).item()) |
| 95 | min_p = torch.clamp(base_min_p * (1 - 0.5 * logits_uncertainty), 0.01, 0.5) |
| 96 | |
| 97 | # Convert tensor values to Python scalars for logging |
| 98 | logging.debug(f"Adaptive sampling params: temp={temperature.item():.3f}, top_p={top_p.item():.3f}, top_k={top_k}, min_p={min_p.item():.3f}") |
| 99 | |
| 100 | samples = [] |
| 101 | for _ in range(n_samples): |
| 102 | sample = _sample(logits, temperature=temperature.item(), top_p=top_p.item(), top_k=top_k, min_p=min_p.item(), generator=generator) |
| 103 | samples.append(sample) |
| 104 | |
| 105 | def score_sample(sample): |
| 106 | sample_flat = sample.flatten().to(torch.long) |
| 107 | one_hot = F.one_hot(sample_flat, logits.shape[-1]) |
| 108 | log_probs = F.log_softmax(logits, dim=-1).view(-1, logits.shape[-1]) |
| 109 | log_prob = torch.sum(log_probs * one_hot) |
| 110 | |
| 111 | confidence_score = ( |
| 112 | (1 - metrics["logits_entropy"]) * 0.1 + |
| 113 | (1 - metrics["attn_entropy"]) * 0.2 + |
| 114 | (1 - metrics["logits_varentropy"]) * 0.3 + |
| 115 | (1 - metrics["attn_varentropy"]) * 0.4 + |
| 116 | metrics["agreement"] * 0.5 + |
| 117 | metrics["interaction_strength"] * 0.6 |
| 118 | ) |
| 119 | return log_prob + confidence_score |
| 120 | |
| 121 | sample_scores = torch.stack([score_sample(sample) for sample in samples]) |
| 122 | best_sample_idx = torch.argmax(sample_scores) |
| 123 | return samples[best_sample_idx] |
| 124 | |
| 125 | def entropy_decode( |
| 126 | model: PreTrainedModel, |
no test coverage detected