Samples from a multinomial distribution using the top-k sampling strategy. Args: logits: A tensor of shape (batch_size, vocab_size) containing the logits. num_samples: The number of samples to draw.
(logits: torch.FloatTensor, num_samples: int)
| 548 | |
| 549 | |
| 550 | def topk_sampling(logits: torch.FloatTensor, num_samples: int): |
| 551 | """ |
| 552 | Samples from a multinomial distribution using the top-k sampling strategy. |
| 553 | |
| 554 | Args: |
| 555 | logits: A tensor of shape (batch_size, vocab_size) containing the logits. |
| 556 | num_samples: The number of samples to draw. |
| 557 | """ |
| 558 | log_prob = F.log_softmax(logits, dim=-1) |
| 559 | topk = torch.topk(log_prob, num_samples, dim=-1) |
| 560 | topk_tokens = topk.indices |
| 561 | topk_log_prob = topk.values |
| 562 | |
| 563 | return topk_tokens, topk_log_prob |
| 564 | |
| 565 | |
| 566 | def nuclear_sampling(logits: torch.FloatTensor, temperature: float, top_p: float = None, top_k: int = None): |
no outgoing calls
no test coverage detected