class_docstr: The dataset class for prompt-tuning or p-tuning pretrained GPT 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
      virtual_prompt_source (Enum): Either VirtualPromptSource.NO_PROMPTS or VirtualPromptSource.PROMPT_ENCODER
      task_templates (dict): Dictionary containing all task template information needed to format prompts. Created in the GPTPromptLearningModel class.
      pseudo_tokens (list[strings]): A list of virtual prompt token placeholders e.g [<prompt_1>, <prompt_2>, ...] up to max num virtual tokens
      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"""
      taskname_ids, input_ids, answer_starts = zip(*batch)
  
      # Pad taskname_ids to be the same length for the prompt encoder
      if (
          self.virtual_prompt_source
          == VirtualPromptSource.PROMPT_ENCODER
      ):
          max_taskname_length = max(len(ids) for ids in taskname_ids)
          taskname_ids = [
              ids
              + [self.pad_token_id] * (max_taskname_length - len(ids))
              for ids in taskname_ids
          ]
          taskname_ids = torch.tensor(taskname_ids)
  
      # Task ids are just used for a look up embeddings for prompt-table
      elif self.virtual_prompt_source == VirtualPromptSource.NO_PROMPT:
          taskname_ids = torch.tensor(taskname_ids)
  
      # Get max sequence length of batch
      batch_max = max(len(ids) for ids in input_ids)
  
      if tp_workers > 1:
          # more 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
      ceil_batch_max = self._ceil_to_nearest(
          batch_max, 8
      )  # @adithyare this padding does not conflict with the tp_workers padding above
      # since tp_workers is always a multiple of 2. the padding to multiple of 8 is to ensure an mem-optimized softmax is used.
      batch_max = ceil_batch_max + 1
      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  # @adithyare I *think* this negatition is done to account for the above 2 lines which removes one item from the input_ids seq.
  
      # Loss mask should align with labels
      loss_mask = loss_mask[:, 1:].contiguous()
  
      # Using causal attention mask for whole input
      batch_size = len(input_ids)
      attention_mask = torch.tril(
          torch.ones((batch_size, batch_max, batch_max))
      ).view(batch_size, 1, batch_max, batch_max)
  
      # Convert attention mask from float to bool
      attention_mask = attention_mask < 0.5
      position_ids = build_position_ids(input_ids)
  
      return (
          input_ids,
          labels,
          loss_mask,
          position_ids,
          attention_mask,
          taskname_ids,
      )
  
imports: import json
  import os
  import pickle
  import torch
  from tqdm.auto import tqdm
  from nemo.collections.nlp.modules.common import VirtualPromptSource
  from nemo.collections.nlp.modules.common.megatron.utils import build_position_ids
  from nemo.core import Dataset
  from nemo.utils import AppState, logging
label_desc: Use image node_id 12 for calling the GPTPromptLearningDataset obj's underlying member method code with example usage: obj.collate_fn(batch, tp_workers) and returns: input_ids,  labels,  loss_mask,  position_ids,  attention_mask,  taskname_ids