code: def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
      """Sample next token from top P probabilities along the last dimension of the input probs tensor.
  
      Args:
          probs (torch.Tensor): Input probabilities with token candidates on the last dimension.
          p (int): The p in “top-p”.
      Returns:
          torch.Tensor: Sampled tokens.
      """
      probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
      probs_sum = torch.cumsum(probs_sort, dim=-1)
      mask = probs_sum - probs_sort > p
      probs_sort *= (~mask).float()
      probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
      next_token = multinomial(probs_sort, num_samples=1)
      next_token = torch.gather(probs_idx, -1, next_token)
      return next_token
  
imports: from concurrent.futures import ProcessPoolExecutor
  from contextlib import contextmanager
  from functools import wraps, lru_cache
  import hashlib
  import json
  import logging
  from pathlib import Path
  import typing
  import flashy
  import flashy.distrib
  import omegaconf
  import torch
  from torch.nn.utils.rnn import pad_sequence
label_desc: Use image node_id 8 for calling a global function with example usage: sample_top_p(probs, p) and returns: next_token