class_docstr: The dataset class for fine tune RETRO models.
  
  Args:
      data (list[strings], list[dicts]): (1) paths to .jsonl or .json files, (2) dict objects corresponding to each input example
      tokenizer (tokenizer): Tokenizer from frozen language model
      task_templates (dict): Dictionary containing all task template information needed to format prompts. Created in the GPTPromptLearningModel class.
      pad_token_id (int): ID of pad token from tokenizer
      max_seq_length (int): maximum sequence length for each dataset examples. Examples will either be truncated to fit this length or dropped if they cannot be truncated.
      min_seq_length (int): min length of each data example in the dataset. Data examples will be dropped if they do not meet the min length requirements. 
      add_bos (bool): Whether to add a beginning of sentence token to each data example
      add_eos (bool): Whether to add an end of sentence token to each data example
      for_train (bool): Whether you're creating a dataset for training or inference
      tokens_to_generate (int): (inference only) Number of tokens to generate during inference
docstr: Prepares input_ids, labels, loss mask, attention_mask, and position ids for global batch 
docstr_tok: ["Prepares","input_ids",",","labels",",","loss","mask",",","attention_mask",",","and","position","ids","for","global","batch"]
code: def collate_fn(self, batch, tp_workers=0):
      """Prepares input_ids, labels, loss mask, attention_mask, and position ids for global batch"""
      input_ids, answer_starts, chunks = zip(*batch)
      # convert chunks into torch tensors
      chunks = torch.tensor(chunks)
  
      # Get max sequence length of batch
      batch_max = max(len(ids) for ids in input_ids)
  
      if tp_workers > 1:
          # make sure the sequence length is multiply of number of tp_workers, needed for sequence parallel.
          resi_padding = (
              tp_workers - (batch_max - 1) % tp_workers
          ) % tp_workers
      else:
          resi_padding = 0
      batch_max += resi_padding
      input_ids, loss_mask = self.pad_batch_and_build_loss_mask(
          input_ids, batch_max, answer_starts
      )
      # Should be a label for every token in batch, label is the next token
      labels = input_ids[:, 1:].contiguous()
      input_ids = input_ids[:, :-1].contiguous()
      batch_max -= 1
  
      # Loss mask should align with labels
      loss_mask = loss_mask[:, 1:].contiguous()
  
      hidden_mask = input_ids != self.pad_token_id
      context_mask = chunks != self.pad_token_id
  
      # Using causal attention mask for whole input
  
      return {
          "tokens": input_ids,
          "labels": labels,
          "tokens_mask": hidden_mask,
          "loss_mask": loss_mask,
          "retrieved_emb_mask": context_mask,
          "retrieved_ids": chunks,
      }
  
file_path: NeMo/nemo/collections/nlp/data/language_modeling/megatron/retro_fine_tune_dataset.py
imports: from random import sample
  import numpy
  import torch
  from nemo.collections.nlp.data.language_modeling.megatron.dataset_utils import get_samples_mapping
  from nemo.collections.nlp.data.language_modeling.text_memmap_dataset import JSONLMemMapDataset
  from nemo.core import Dataset
  from nemo.utils import logging
label_desc: Use image node_id 6 for calling the RetroQAFineTuneDataset obj's underlying member method code with example usage: obj.collate_fn(batch, tp_workers) and returns: dict