code: def get_batch(data_iterator):
      """Generate a batch"""
      args = get_args()
      retro_args = get_retro_args()
      tokenizer = get_tokenizer()
  
      # Items and their type.
      keys = ["text"]
      if args.retro_add_retriever:
          keys.append("neighbor_tokens")
      datatype = torch.int64
  
      # Broadcast data.
      if data_iterator is not None:
          data = next(data_iterator)
      else:
          data = None
  
      data_b = tensor_parallel.broadcast_data(keys, data, datatype)
  
      # Unpack.
      tokens_ = data_b["text"].long()
      labels = tokens_[:, 1:].contiguous()
      tokens = tokens_[:, :-1].contiguous()
  
      if args.retro_add_retriever:
          # note: [bs * l * k, r]
          # note: 2x == neighbor, continuation
          neighbor_tokens = (
              data_b["neighbor_tokens"]
              .view(-1, retro_args.retro_gpt_retrieved_length)
              .long()
          )
  
      # Get the masks and postition ids.
      (
          attention_mask,
          loss_mask,
          position_ids,
      ) = get_ltor_masks_and_position_ids(
          tokens,
          tokenizer.eod,
          args.reset_position_ids,
          args.reset_attention_mask,
          args.eod_mask_loss,
      )
  
      if args.retro_add_retriever:
          _, _, neighbor_position_ids = get_ltor_masks_and_position_ids(
              neighbor_tokens,
              tokenizer.eod,
              args.reset_position_ids,
              args.reset_attention_mask,
              args.eod_mask_loss,
          )
          neighbor_attention_mask = None
          return (
              tokens,
              labels,
              loss_mask,
              attention_mask,
              position_ids,
              neighbor_tokens,
              neighbor_attention_mask,
              neighbor_position_ids,
          )
  
      else:
          return tokens, labels, loss_mask, attention_mask, position_ids
  
imports: from functools import partial
  import torch
  from megatron import get_args, get_retro_args
  from megatron import get_timers
  from megatron import get_tokenizer
  from megatron import print_rank_0
  from megatron.arguments import core_transformer_config_from_args
  from megatron.core import tensor_parallel
  from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
  from megatron.core.datasets.gpt_dataset import GPTDataset
  from megatron.core.enums import ModelType
  from megatron.core.models.retro import get_retro_decoder_block_spec, RetroModel
  from megatron.training import pretrain
  from megatron.utils import get_ltor_masks_and_position_ids
  from tools.retro.query.chunk_dataset import train_valid_test_datasets_provider
  from tools.retro.query.retro_dataset import get_retro_datasets
  from pretrain_gpt import loss_func, model_provider
label_desc: Use image node_id 3 for calling a global function with example usage: get_batch(data_iterator) and returns: tokens,  labels,  loss_mask,  attention_mask,  position_ids,  neighbor_tokens,  neighbor_attention_mask,  neighbor_position_ids, tokens,  labels,  loss_mask,  attention_mask,  position_ids