def topk_sampling(logits, top_k=50, top_p=1.0, temperature=1.0):
    """
    Perform top-k and top-p sampling on logits.

    Args:
        logits (torch.Tensor): The logits to sample from.
        top_k (int, optional): The number of highest probability tokens to keep for top-k filtering.
                               Must be a positive integer. Defaults to 50.
        top_p (float, optional): The cumulative probability threshold for nucleus sampling.
                                 Must be between 0 and 1. Defaults to 1.0.
        temperature (float, optional): The scaling factor to adjust the logits distribution.
                                       Must be strictly positive. Defaults to 1.0.

    Returns:
        torch.Tensor: The sampled token.
    """

    # Adjust logits using temperature
    if temperature != 1.0:
        logits = logits / temperature

    # Top-p/top-k filtering
    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)

    # Sample from the filtered distribution
    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
    return token
Perform top-k and top-p sampling on logits.

Args:
logits (torch.Tensor): The logits to sample from.
top_k (int, optional): The number of highest probability tokens to keep for top-k filtering.
Must be a positive integer. Defaults to 50.
top_p (float, optional): The cumulative probability threshold for nucleus sampling.
Must be between 0 and 1. Defaults to 1.0.
temperature (float, optional): The scaling factor to adjust the logits distribution.
Must be strictly positive. Defaults to 1.0.

Returns:
torch.Tensor: The sampled token.
utils/topk_sampling.py