(
logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None, top_p: float = 1.0
)
| 34 | |
| 35 | |
| 36 | def sample( |
| 37 | logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None, top_p: float = 1.0 |
| 38 | ) -> torch.Tensor: |
| 39 | if top_p < 0.0 or top_p > 1.0: |
| 40 | raise ValueError(f"top_p must be in [0, 1], got {top_p}") |
| 41 | logits = logits[0, -1] |
| 42 | # optionally crop the logits to only the top k options |
| 43 | if top_k is not None: |
| 44 | v, i = torch.topk(logits, min(top_k, logits.size(-1))) |
| 45 | # do not use `torch.where` as in nanogpt because it will repeat top-k collisions |
| 46 | logits = torch.full_like(logits, float("-inf")).scatter_(-1, i, v) |
| 47 | # optionally scale the logits and sample from a probability distribution |
| 48 | if temperature > 0.0 or top_p > 0.0: |
| 49 | if temperature > 0.0: |
| 50 | logits = logits / temperature |
| 51 | # optionally crop the logits to smallest set of logits with a cumulative probability above top_p |
| 52 | if top_p < 1.0: |
| 53 | logits = sample_top_p(logits, top_p) |
| 54 | probs = torch.nn.functional.softmax(logits, dim=-1) |
| 55 | return multinomial_num_samples_1(probs) |
| 56 | return torch.argmax(logits, dim=-1, keepdim=True) |
| 57 | |
| 58 | |
| 59 | def next_token( |
no test coverage detected