Method _sample_next_token

python/TIDE/runtime.py:316–352

Sample a single next token from logits [B, vocab].

(
        logits: torch.Tensor,
        temperature: float,
        top_k: int,
        top_p: float,
    )

    @staticmethod
    def _sample_next_token(
        logits: torch.Tensor,
        temperature: float,
        top_k: int,
        top_p: float,
    ) -> torch.Tensor:
        """Sample a single next token from logits [B, vocab]."""
        next_logits = logits.clone()

        if temperature > 0:
            next_logits = next_logits / temperature

            if top_k > 0:
                topk_vals, _ = torch.topk(next_logits, top_k)
                next_logits[next_logits < topk_vals[:, -1:]] = float("-inf")

            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(
                    next_logits, descending=True
                )
                cumulative_probs = (
                    torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
                )
                remove_mask = cumulative_probs > top_p
                remove_mask[:, 1:] = remove_mask[:, :-1].clone()
                remove_mask[:, 0] = False
                sorted_logits[remove_mask] = float("-inf")
                next_logits = sorted_logits.scatter(
                    1, sorted_indices, sorted_logits
                )

            probs = torch.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
        else:
            next_token = next_logits.argmax(dim=-1, keepdim=True)

        return next_token

    @torch.no_grad()
    def generate(