(logits: torch.Tensor, temperature=0.666, top_p=0.90, top_k=27, min_p: float = 0.0, generator: torch.Generator = None)
| 58 | } |
| 59 | |
| 60 | def _sample(logits: torch.Tensor, temperature=0.666, top_p=0.90, top_k=27, min_p: float = 0.0, generator: torch.Generator = None) -> torch.Tensor: |
| 61 | bsz = logits.shape[0] |
| 62 | logit = logits[:, -1] |
| 63 | probs = F.softmax(logit / temperature, dim=-1) |
| 64 | |
| 65 | if min_p > 0.0: |
| 66 | p_max = torch.max(probs, dim=-1, keepdim=True).values |
| 67 | indices_to_remove = probs < (min_p * p_max) |
| 68 | logit = torch.where(indices_to_remove, torch.full_like(logit, float('-inf')), logit) |
| 69 | |
| 70 | top_k_probs, top_k_indices = torch.topk(probs, k=min(top_k, probs.shape[-1])) |
| 71 | probs_sort = torch.flip(top_k_probs, dims=[-1]) |
| 72 | probs_idx = torch.flip(top_k_indices, dims=[-1]) |
| 73 | probs_sum = torch.cumsum(probs_sort, dim=-1) |
| 74 | mask = torch.where(probs_sum - probs_sort > top_p, torch.tensor(1.0, device=device), torch.tensor(0.0, device=device)) |
| 75 | probs_sort = probs_sort * (1 - mask) |
| 76 | probs_sort = probs_sort / torch.sum(probs_sort, dim=-1, keepdim=True) |
| 77 | next_token = torch.multinomial(probs_sort, 1, generator=generator) |
| 78 | next_token_g = torch.gather(probs_idx, -1, next_token.reshape(bsz, 1).to(torch.int64)) |
| 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, |
no outgoing calls
no test coverage detected