Top-p (Nucleus) Sampling with Logits or Probabilities

Requirements Accept a one-dimensional NumPy or PyTorch tensor containing either logits or vocabulary probabilities for V tokens, together with a boolean input_is_logits flag indicating which representation is provided, and a cutoff $$p \in (0, 1]$$. When the tensor contains logits, obtain probabilities using a numerically stable softmax based on log-sum-exp. Order the probabilities from largest to smallest and add them cumulatively; select the shortest leading segment whose…