from __future__ import annotations

import torch


def sample_token(
    logits: torch.Tensor,
    *,
    do_sample: bool = False,
    top_k: int = 50,
    top_p: float = 1.0,
    temperature: float = 0.9,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    if not do_sample:
        return logits.argmax(dim=-1)
    scores = logits.float() / temperature
    if 0 < top_k < scores.shape[-1]:
        threshold = scores.topk(top_k, dim=-1).values[..., -1, None]
        scores = scores.masked_fill(scores < threshold, -torch.inf)
    if top_p < 1.0:
        sorted_scores, sorted_indices = scores.sort(dim=-1, descending=True)
        remove = sorted_scores.softmax(dim=-1).cumsum(dim=-1) > top_p
        remove[..., 1:] = remove[..., :-1].clone()
        remove[..., 0] = False
        sorted_scores = sorted_scores.masked_fill(remove, -torch.inf)
        scores = torch.full_like(scores, -torch.inf).scatter(
            -1, sorted_indices, sorted_scores
        )
    return torch.multinomial(
        scores.softmax(dim=-1), num_samples=1, generator=generator
    ).squeeze(-1)