code: def score_and_return_on_first_stage(model, tokens, lengths):
      """Function for just scoring.
      Arguments:
          model: no interleaving is supported.
          tokens: prompt tokens extended to be of size [b, max_prompt_length]
          lengths: original prompt length, size: [b]
      Note: Outside of model, other parameters only need to be available on
            rank 0.
      Outputs:
          output_log_probs: log probability of the selected tokens. size: [b, s]
      """
  
      args = get_args()
  
      batch_size = tokens.size(0)
      max_prompt_length = lengths.max().item()
      assert max_prompt_length == tokens.size(1)
  
      if max_prompt_length > args.max_position_embeddings:
          raise ValueError(
              "Length of prompt + tokens_to_generate longer than allowed"
          )
  
      if max_prompt_length * batch_size > args.max_tokens_to_oom:
          raise ValueError(
              "Too many tokens.  "
              + str(max_prompt_length * batch_size)
              + " is greater than "
              + str(args.max_tokens_to_oom)
          )
  
      # forward step.
      forward_step = ForwardStep(model, batch_size, max_prompt_length)
  
      # ===================
      # Pre-allocate memory
      # ===================
  
      # Log probability of the sequence (prompt + generated tokens).
      output_log_probs = None
      output_log_probs_size = (batch_size, max_prompt_length - 1)
  
      if mpu.is_pipeline_last_stage():
          output_log_probs = torch.empty(
              output_log_probs_size,
              dtype=torch.float32,
              device=torch.cuda.current_device(),
          )
  
      # =============
      # Run infernece
      # =============
      with torch.no_grad():
          (
              attention_mask,
              position_ids,
          ) = _build_attention_mask_and_position_ids(tokens)
  
          # logits will be meanigful only in the last pipeline stage.
          logits = forward_step(tokens, position_ids, attention_mask)
  
          if mpu.is_pipeline_last_stage():
              # Always the last stage should have an output.
              assert logits is not None
              log_probs = F.log_softmax(logits, dim=2)
  
              # Pick the tokens that we need to get the log
              # probabilities for. Note that next input token is
              # the token which we selected in the current logits,
              # so shift by 1.
              indices = torch.unsqueeze(tokens[:, 1:], 2)
              output_log_probs = torch.gather(
                  log_probs, 2, indices
              ).squeeze(2)
  
      # ======================================
      # Broadcast to the first pipeline stage.
      # ======================================
      output_log_probs = broadcast_from_last_to_first_pipeline_stage(
          output_log_probs_size, torch.float32, output_log_probs
      )
  
      return tokens, lengths, output_log_probs, logits
  
file_path: megatron-lm/megatron/text_generation/generation.py
imports: import torch
  import torch.nn.functional
  from megatron import get_args, get_tokenizer
  from megatron.core import mpu
  from megatron.utils import get_ltor_masks_and_position_ids
  from .communication import copy_from_last_to_first_pipeline_stage, broadcast_from_last_pipeline_stage, broadcast_from_last_to_first_pipeline_stage
  from .forward_step import ForwardStep
  from .sampling import sample
  from .beam_utils import BeamHypotheses
label_desc: Use image node_id 1 for calling a global function with example usage: score_and_return_on_first_stage(model, tokens, lengths) and returns: tokens,  lengths,  output_log_probs,  logits