code: def sample_token_greedy(logits):
      """
      Greedy sampling. Returns the token with the highest probability, and corresponding log_prob.
  
      Args:
          logits: [batch_size, vocab_size] - unnormalized log probabilities of the next token
  
      Returns:
          log_probs: [batch_size] - log probabilities of the sampled tokens
          token_ids: [batch_size] - sampled token ids
      """
      log_probs, token_ids = torch.max(
          torch.nn.functional.log_softmax(logits, dim=-1), dim=-1
      )
  
      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 18 for calling a global function with example usage: sample_token_greedy(logits) and returns: log_probs,  token_ids