code: def batch_top_tokens(
      top_n_tokens: List[int],
      top_n_tokens_tensor: torch.Tensor,
      logprobs: torch.Tensor,
  ) -> Tuple[List[List[int]], List[List[float]]]:
      """Find the top n most likely tokens for a batch of generations.
  
      When multiple tokens have equal probabilities and they don't all fit, the
      remaining tokens are also returned.
      """
      max_top_n = max(top_n_tokens)
      # Early exit when top_n_tokens is not used
      if max_top_n == 0:
          return [[]] * len(top_n_tokens), [[]] * len(top_n_tokens)
  
      # Ensure top_n doesn't exceed vocab size
      top_n_tokens = [
          min(tok, logprobs.size(-1)) for tok in top_n_tokens
      ]
  
      # Parallel kthvalue adapted from https://discuss.pytorch.org/t/how-to-efficiently-get-the-k-th-largest-values-in-parallel/160529/2
      # Sorted topk is faster than torch.sort() since we only need a small subset
      sorted_top_k = torch.topk(
          logprobs, k=max_top_n, dim=1, sorted=True
      ).values
      nth_highest = torch.gather(
          sorted_top_k,
          1,
          (top_n_tokens_tensor - 1).clip(min=0).unsqueeze(1),
      )
      nth_highest[nth_highest == -float("inf")] = torch.finfo(
          logprobs.dtype
      ).min
  
      # Find the new "fuzzy" top n values
      top_n_indices = (logprobs >= nth_highest).nonzero()
      _, top_n_ishes = torch.unique_consecutive(
          top_n_indices[:, 0], return_counts=True
      )
  
      k = 1 if top_n_ishes.numel() == 0 else top_n_ishes.max()
      # Take a new topk for these new max n values
      top_k = torch.topk(logprobs, k=k, dim=1, sorted=True)
  
      top_n_ishes = top_n_ishes.tolist()
      top_indices = top_k.indices.tolist()
      top_values = top_k.values.tolist()
  
      return (
          [
              idxs[:n] if req_n > 0 else []
              for idxs, n, req_n in zip(
                  top_indices, top_n_ishes, top_n_tokens
              )
          ],
          [
              vals[:n] if req_n > 0 else []
              for vals, n, req_n in zip(
                  top_values, top_n_ishes, top_n_tokens
              )
          ],
      )
  
imports: import re
  from typing import Callable, List, Optional, Tuple
  import torch
  from text_generation_server.pb import generate_pb2
  from text_generation_server.pb.generate_pb2 import FinishReason
  from text_generation_server.utils.logits_process import HeterogeneousProcessorWrapper, HeterogeneousRepetitionPenaltyLogitsProcessor, HeterogeneousTemperatureLogitsWarper, HeterogeneousTopKLogitsWarper, HeterogeneousTopPLogitsWarper, HeterogeneousTypicalLogitsWarper, static_warper
  from text_generation_server.utils.watermark import WatermarkLogitsProcessor
  from transformers import PreTrainedTokenizerBase, RepetitionPenaltyLogitsProcessor
label_desc: Use image node_id 2 for calling a global function with example usage: batch_top_tokens(top_n_tokens, top_n_tokens_tensor, logprobs) and returns: unknown,  unknown, unknown,  unknown