(logits: torch.FloatTensor, temperature: float, top_p: float = None, top_k: int = None)
| 564 | |
| 565 | |
| 566 | def nuclear_sampling(logits: torch.FloatTensor, temperature: float, top_p: float = None, top_k: int = None): |
| 567 | orig_log_probs = F.log_softmax(logits, dim=-1) |
| 568 | logits /= temperature |
| 569 | logits = top_k_logits(logits, top_k, top_p) |
| 570 | log_probs = F.softmax(logits, dim=-1) |
| 571 | tokens = torch.multinomial(log_probs, num_samples=1).view(-1) |
| 572 | |
| 573 | indices = tokens.view(-1, 1) |
| 574 | new_scores = orig_log_probs.gather(1, indices).view(-1) |
| 575 | |
| 576 | return tokens, new_scores |
| 577 | |
| 578 | |
| 579 | def sample_topk_tokens(model, |
no test coverage detected