code: def get_samples_mapping_(
      indexed_dataset,
      data_prefix,
      num_epochs,
      max_num_samples,
      max_seq_length,
      short_seq_prob,
      seed,
      name,
      binary_head,
  ):
      logger = get_dist_logger()
      if not num_epochs:
          if not max_num_samples:
              raise ValueError(
                  "Need to specify either max_num_samples "
                  "or num_epochs"
              )
          num_epochs = np.iinfo(np.int32).max - 1
      if not max_num_samples:
          max_num_samples = np.iinfo(np.int64).max - 1
  
      # Filename of the index mapping
      indexmap_filename = data_prefix
      indexmap_filename += "_{}_indexmap".format(name)
      if num_epochs != (np.iinfo(np.int32).max - 1):
          indexmap_filename += "_{}ep".format(num_epochs)
      if max_num_samples != (np.iinfo(np.int64).max - 1):
          indexmap_filename += "_{}mns".format(max_num_samples)
      indexmap_filename += "_{}msl".format(max_seq_length)
      indexmap_filename += "_{:0.2f}ssp".format(short_seq_prob)
      indexmap_filename += "_{}s".format(seed)
      indexmap_filename += ".npy"
  
      # Build the indexed mapping if not exist.
      if torch.distributed.get_rank() == 0 and not os.path.isfile(
          indexmap_filename
      ):
          print(
              " > WARNING: could not find index map file {}, building "
              "the indices on rank 0 ...".format(indexmap_filename)
          )
  
          # Make sure the types match the helpers input types.
          assert indexed_dataset.doc_idx.dtype == np.int64
          assert indexed_dataset.sizes.dtype == np.int32
  
          # Build samples mapping
          verbose = torch.distributed.get_rank() == 0
          start_time = time.time()
          logger.info(
              "\n > building samples index mapping for {} ...".format(
                  name
              ),
              ranks=[0],
          )
          # First compile and then import.
          samples_mapping = helpers.build_mapping(
              indexed_dataset.doc_idx,
              indexed_dataset.sizes,
              num_epochs,
              max_num_samples,
              max_seq_length,
              short_seq_prob,
              seed,
              verbose,
              2 if binary_head else 1,
          )
          logger.info(
              "\n > done building samples index maping", ranks=[0]
          )
          np.save(indexmap_filename, samples_mapping, allow_pickle=True)
          logger.info(
              "\n > saved the index mapping in {}".format(
                  indexmap_filename
              ),
              ranks=[0],
          )
          # Make sure all the ranks have built the mapping
          logger.info(
              "\n > elapsed time to build and save samples mapping "
              "(seconds): {:4f}".format(time.time() - start_time),
              ranks=[0],
          )
      # This should be a barrier but nccl barrier assumes
      # device_index=rank which is not the case for model
      # parallel case
      counts = torch.cuda.LongTensor([1])
      torch.distributed.all_reduce(
          counts, group=gpc.get_group(ParallelMode.DATA)
      )
      if gpc.is_initialized(ParallelMode.PIPELINE):
          torch.distributed.all_reduce(
              counts, group=gpc.get_group(ParallelMode.PIPELINE)
          )
      assert counts[0].item() == (
          torch.distributed.get_world_size()
          // torch.distributed.get_world_size(
              group=gpc.get_group(ParallelMode.SEQUENCE)
          )
      )
  
      # Load indexed dataset.
      start_time = time.time()
      samples_mapping = np.load(
          indexmap_filename, allow_pickle=True, mmap_mode="r"
      )
      logger.info(
          "\n > loading indexed mapping from {}".format(
              indexmap_filename
          )
          + "\n    loaded indexed file in {:3.3f} seconds".format(
              time.time() - start_time
          )
          + "\n    total number of samples: {}".format(
              samples_mapping.shape[0]
          ),
          ranks=[0],
      )
  
      return samples_mapping
  
file_path: ColossalAI/examples/tutorial/sequence_parallel/data/datasets/bert_dataset.py
imports: import os
  import time
  import numpy
  import torch
  from torch.utils.data import Dataset
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.logging import get_dist_logger
  from ..tokenizer import get_tokenizer
  from .dataset_utils import create_masked_lm_predictions, create_tokens_and_tokentypes, get_a_and_b_segments, pad_and_convert_to_numpy, truncate_segments
label_desc: Use image node_id 1 for calling a global function with example usage: get_samples_mapping_(indexed_dataset, data_prefix, num_epochs, max_num_samples, max_seq_length, short_seq_prob, seed, name, binary_head) and returns: samples_mapping