import torch


def top_p_filter(
    logits: torch.Tensor,
    top_p: float,
    min_tokens_to_keep: int,
    is_probs: bool = False
) -> torch.Tensor:
    """Helper function for nucleus sampling decoding, aka. top-p decoding.

    :param logits:
    :param top_p:
    :param min_tokens_to_keep:
    :param is_probs:
    :return:
    """
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)

    sorted_indices_to_remove = cumulative_probs > top_p
    if min_tokens_to_keep > 1:
        # keep at least min tokens
        sorted_indices_to_remove[..., :min_tokens_to_keep - 1] = 0
    # Shift the indices to the right to keep also the first token above the threshold
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0

    # scatter sorted tensors to original indexing
    indices_to_remove = sorted_indices_to_remove.scatter(
        1, sorted_indices, sorted_indices_to_remove
    )
    if is_probs:
        scores = logits.masked_fill(indices_to_remove, 0.0)
    else:
        scores = logits.masked_fill(indices_to_remove, -float("Inf"))
    return scores