code: def get_sampling_token_fn(
      sampling_method: str, sampling_kwargs: dict
  ) -> Tuple[Callable, dict]:
      """
      Specifies the sampling function that takes in a tensor of logits [batch_size, vocab_size] and returns a tuple
      (tensor of log_probs [batch_size], tensor of sampled from logits [batch_size]).
      If the beam search is enabled, the sampling function returns tensors [batch_size, beam_size]
  
      Args:
          sampling_method: the sampling method to use in the decode steps. Currently supported methods are
                            "beam-search"/"greedy"/"topkp"
          sampling_kwargs: dict with arguments to be passed to the sampling function.
                           For sampling method 'beam-search', the following kwargs are supported:
                           beam_size - int, number of the best sequences at each decode iteration to be left per target
                           beam_alpha - int, the parameter of length penalty applied to predicted sequences
                           keep_only_best_tokens - used in the beam search, boolean flag if to output only best sequence
                                                   of predicted tokens (True) or beam_size predictions per target
                           return_scores - used in the beam search, boolean flag if to return scores at the top of
                                           predictions and logits
  
      Returns:
          sample_token_fn: the sampling function
          default_sampling_kwargs: sampling_kwargs augmented with default sampling kwargs
      """
      all_default_sampling_kwargs = {
          "greedy-search": {},
          "topkp-sampling": {
              "top_k": 0,
              "top_p": 0.0,
              "temperature": 1.0,
          },
          "beam-search": {
              "beam_size": 1,
              "beam_alpha": 0.0,
              "keep_only_best_tokens": False,
              "return_scores": False,
          },
      }
  
      # update default sampling kwargs with user provided kwargs
      default_sampling_kwargs = all_default_sampling_kwargs[
          sampling_method
      ].copy()
      default_sampling_kwargs.update(sampling_kwargs)
      # sampling_kwargs = default_sampling_kwargs
  
      if sampling_method == "greedy-search":
          sampling_token_fn = sample_token_greedy
  
      elif sampling_method == "topkp-sampling":
          top_k = default_sampling_kwargs["top_k"]
          top_p = default_sampling_kwargs["top_p"]
          temperature = default_sampling_kwargs["temperature"]
          sampling_token_fn = partial(
              sample_token_topk,
              top_k=top_k,
              top_p=top_p,
              temperature=temperature,
          )
  
      elif sampling_method == "beam-search":
          beam_size = default_sampling_kwargs["beam_size"]
          sampling_token_fn = partial(
              sample_token_topk_beam_search, beam_size=beam_size
          )
  
      else:
          raise ValueError(
              f"Invalid sampling method {sampling_method}. "
              f"Supported sampling methods are {all_default_sampling_kwargs.keys()}"
          )
  
      return sampling_token_fn, default_sampling_kwargs
  
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 22 for calling a global function with example usage: get_sampling_token_fn(sampling_method, sampling_kwargs) and returns: sampling_token_fn,  default_sampling_kwargs