code: def sample_top_p(probs, p):
      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] = 0.0
      probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
      next_token = torch.multinomial(probs_sort, num_samples=1)
      next_token = torch.gather(probs_idx, -1, next_token)
      return next_token
  
imports: import json
  import os
  import sys
  import time
  from pathlib import Path
  from typing import List, Literal, Optional, Tuple, TypedDict
  import torch
  import torch.nn.functional
  from fairscale.nn.model_parallel.initialize import get_model_parallel_rank, initialize_model_parallel, model_parallel_is_initialized
  from llama.model import ModelArgs, Transformer
  from llama.tokenizer import Tokenizer
label_desc: Use image node_id 1 for calling a global function with example usage: sample_top_p(probs, p) and returns: next_token