code: def sample_token_topk_beam_search(
      logits: torch.Tensor,
      beam_size: int = 1,
      dim: int = -1,
      log_softmax: bool = True,
  ):
      """
      Beam search selection of top K predictions per target (dim). Returns the beam_size tokens ids with the highest
      probability and the corresponding log_prob per target
  
      Args:
          logits: [batch_size, vocab_size] or [batch_size, vocab_size] - unnormalized log probabilities of the next token,
          beam_size: int > 1 - number of tokens to return with the highest probability per target
          dim: int - dim of log_softmax and topk selection
          log_softmax: bool - if to calculate log softmax  for log probabilities
  
  
      Returns:
          log_probs: [batch_size, beam_size] - log probabilities of the sampled tokens
          token_ids: [batch_size, beam_size] - sampled token ids
      """
      if log_softmax:
          log_probs = torch.nn.functional.log_softmax(logits, dim=dim)
      else:
          log_probs = logits
      # get top candidates for each item in batch
      log_probs, token_ids = torch.topk(log_probs, beam_size, dim=dim)
  
      return log_probs, token_ids
  
imports: import os
  import pickle
  import re
  from collections.abc import Iterable
  from functools import partial
  from typing import Callable, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from nemo.collections.common.tokenizers.tabular_tokenizer import TabularTokenizer
  from nemo.collections.nlp.modules.common.megatron.utils import get_ltor_masks_and_position_ids
  from nemo.collections.nlp.modules.common.text_generation_strategy import model_inference_strategy_dispatcher
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, OutputType, SamplingParam
  from nemo.utils import AppState
label_desc: Use image node_id 20 for calling a global function with example usage: sample_token_topk_beam_search(logits, beam_size, dim, log_softmax) and returns: log_probs,  token_ids