| 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) |