Return the set of token indices allowed after top-k then top-p filtering.
(
logits: torch.Tensor, # [V], float32
temperature: torch.Tensor,
top_k: int,
top_p: float,
)
| 179 | |
| 180 | |
| 181 | def reference_top_k_top_p( |
| 182 | logits: torch.Tensor, # [V], float32 |
| 183 | temperature: torch.Tensor, |
| 184 | top_k: int, |
| 185 | top_p: float, |
| 186 | ) -> torch.Tensor: |
| 187 | """Return the set of token indices allowed after top-k then top-p filtering.""" |
| 188 | scaled = logits.float() / temperature |
| 189 | topk_values, topk_indices = scaled.topk(top_k) |
| 190 | probs_topk = topk_values.softmax(dim=-1) |
| 191 | sorted_probs, sorted_order = probs_topk.sort(descending=True) |
| 192 | cumsum = sorted_probs.cumsum(dim=0) |
| 193 | mask = cumsum - sorted_probs < top_p |
| 194 | return topk_indices[sorted_order[mask]] |
| 195 | |
| 196 | |
| 197 | @pytest.mark.parametrize("n_hidden_states", [1, 2]) |