code: def _build_index_mappings(
      name,
      data_prefix,
      documents,
      sizes,
      num_samples,
      seq_length,
      seed,
      index_mapping_dir: str = None,
      drop_last: bool = True,
      add_extra_token: int = 1,
      shuffle_documents: bool = True,
      exchange_indices_distributed: bool = False,
  ):
      """Build doc-idx, sample-idx, and shuffle-idx.
      doc-idx: is an array (ordered) of documents to be used in training.
      sample-idx: is the start document index and document offset for each
         training sample.
      shuffle-idx: maps the sample index into a random index into sample-idx.
      """
      # Number of tokens in each epoch and number of required epochs.
      tokens_per_epoch = _num_tokens(documents, sizes)
      num_epochs = _num_epochs(
          tokens_per_epoch, seq_length, num_samples, add_extra_token
      )
      # rng state
      np_rng = np.random.RandomState(seed=seed)
  
      # Filename of the index mappings.
      if index_mapping_dir is not None:
          _filename = os.path.join(
              index_mapping_dir, os.path.basename(data_prefix)
          )
      else:
          _filename = data_prefix
      _filename += "_{}_indexmap".format(name)
      _filename += "_{}ns".format(num_samples)
      _filename += "_{}sl".format(seq_length)
      _filename += "_{}s".format(seed)
      doc_idx_filename = _filename + "_doc_idx.npy"
      sample_idx_filename = _filename + "_sample_idx.npy"
      shuffle_idx_filename = _filename + "_shuffle_idx.npy"
  
      # Build the indexed mapping if not exist.
      if torch.distributed.get_rank() == 0:
          using_cached_indices = True
          if (
              (not os.path.isfile(doc_idx_filename))
              or (not os.path.isfile(sample_idx_filename))
              or (not os.path.isfile(shuffle_idx_filename))
          ):
              using_cached_indices = False
              logging.info(
                  " > WARNING: could not find index map files, building "
                  "the indices on rank 0 ..."
              )
  
              # For the last epoch, decide whether include the entire epoch
              # in the global shuffle or not.
  
              # If we need only one epoch, then separating last epoch  does
              # not mean anything.
              if num_epochs == 1:
                  separate_last_epoch = False
                  print(
                      " > only one epoch required, setting "
                      "separate_last_epoch to False",
                      flush=True,
                  )
  
              else:
                  # Get the number of samples for the last epoch
                  num_samples_from_epochs_minus_one = (
                      (num_epochs - 1) * tokens_per_epoch
                      - add_extra_token
                  ) // seq_length
                  last_epoch_num_samples = (
                      num_samples - num_samples_from_epochs_minus_one
                  )
                  assert (
                      last_epoch_num_samples >= 0
                  ), "last epoch number of samples should be non-negative."
                  num_samples_per_epoch = (
                      tokens_per_epoch - add_extra_token
                  ) // seq_length
                  assert last_epoch_num_samples <= (
                      num_samples_per_epoch + 1
                  ), "last epoch number of samples exceeded max value."
                  # If we have less than 80% of the samples for the last epoch,
                  # seperate out the epoch and treat it differently.
                  # Note: the 80% number is just based on common sense and can
                  # be adjusted if needed.
                  separate_last_epoch = last_epoch_num_samples < int(
                      0.80 * num_samples_per_epoch
                  )
                  if separate_last_epoch:
                      string = (
                          " > last epoch number of samples ({}) is smaller "
                          "than 80% of number of samples per epoch ({}), "
                          "setting separate_last_epoch to True"
                      )
                  else:
                      string = (
                          " > last epoch number of samples ({}) is larger "
                          "than 80% of number of samples per epoch ({}), "
                          "setting separate_last_epoch to False"
                      )
                  print(
                      string.format(
                          last_epoch_num_samples, num_samples_per_epoch
                      ),
                      flush=True,
                  )
  
              # doc-idx.
              start_time = time.time()
              doc_idx = _build_doc_idx(
                  documents,
                  num_epochs,
                  np_rng,
                  separate_last_epoch,
                  shuffle_documents,
              )
              np.save(doc_idx_filename, doc_idx, allow_pickle=True)
              logging.info(
                  " > elasped time to build and save doc-idx mapping "
                  "(seconds): {:4f}".format(time.time() - start_time)
              )
              # sample-idx.
              start_time = time.time()
              # Use C++ implementation for speed.
              # First compile and then import.
              assert doc_idx.dtype == np.int32
              assert sizes.dtype == np.int32
              try:
                  from nemo.collections.nlp.data.language_modeling.megatron.dataset_utils import (
                      compile_helper,
                  )
  
                  compile_helper()
                  from nemo.collections.nlp.data.language_modeling.megatron import (
                      helpers,
                  )
              except ImportError:
                  raise ImportError(
                      f"Could not compile megatron dataset C++ helper functions and therefore cannot import helpers python file."
                  )
  
              sample_idx = helpers.build_sample_idx(
                  sizes,
                  doc_idx,
                  seq_length,
                  num_epochs,
                  tokens_per_epoch,
                  drop_last,
                  add_extra_token,
              )
              # sample_idx = _build_sample_idx(sizes, doc_idx, seq_length,
              #                              num_epochs, tokens_per_epoch, drop_last, add_extra_token)
              np.save(
                  sample_idx_filename, sample_idx, allow_pickle=True
              )
              logging.info(
                  " > elasped time to build and save sample-idx mapping "
                  "(seconds): {:4f}".format(time.time() - start_time)
              )
              # shuffle-idx.
              start_time = time.time()
              # -1 is due to data structure used to retieve the index:
              #    sample i --> [sample_idx[i], sample_idx[i+1])
              if separate_last_epoch:
                  num_samples_ = num_samples_from_epochs_minus_one
              else:
                  num_samples_ = sample_idx.shape[0] - 1
              shuffle_idx = _build_shuffle_idx(
                  num_samples_, sample_idx.shape[0] - 1, np_rng
              )
              np.save(
                  shuffle_idx_filename, shuffle_idx, allow_pickle=True
              )
              logging.info(
                  " > elasped time to build and save shuffle-idx mapping"
                  " (seconds): {:4f}".format(time.time() - start_time)
              )
  
      torch.distributed.barrier()
      counts = torch.cuda.LongTensor([1])
      torch.distributed.all_reduce(
          counts,
          group=parallel_state.get_data_parallel_group(
              with_context_parallel=True
          ),
      )
      torch.distributed.all_reduce(
          counts,
          group=parallel_state.get_pipeline_model_parallel_group(),
      )
      assert counts[0].item() == (
          torch.distributed.get_world_size()
          // torch.distributed.get_world_size(
              group=parallel_state.get_tensor_model_parallel_group()
          )
      )
  
      if not exchange_indices_distributed or (
          torch.distributed.get_rank() == 0 and using_cached_indices
      ):
          # Load mappings.
          start_time = time.time()
          logging.info(
              " > loading doc-idx mapping from {}".format(
                  doc_idx_filename
              )
          )
          doc_idx = np.load(
              doc_idx_filename, allow_pickle=True, mmap_mode="r"
          )
          logging.info(
              " > loading sample-idx mapping from {}".format(
                  sample_idx_filename
              )
          )
          sample_idx = np.load(
              sample_idx_filename, allow_pickle=True, mmap_mode="r"
          )
          logging.info(
              " > loading shuffle-idx mapping from {}".format(
                  shuffle_idx_filename
              )
          )
          shuffle_idx = np.load(
              shuffle_idx_filename, allow_pickle=True, mmap_mode="r"
          )
          logging.info(
              "    loaded indexed file in {:3.3f} seconds".format(
                  time.time() - start_time
              )
          )
          logging.info(
              "    total number of samples: {}".format(
                  sample_idx.shape[0]
              )
          )
          logging.info(
              "    total number of epochs: {}".format(num_epochs)
          )
  
      if exchange_indices_distributed:
          if torch.distributed.get_rank() == 0:
              indices = [(doc_idx, sample_idx, shuffle_idx)]
          else:
              indices = [None]
          torch.distributed.broadcast_object_list(indices)
          doc_idx, sample_idx, shuffle_idx = indices[0]
  
      return doc_idx, sample_idx, shuffle_idx
  
imports: import os
  import time
  import numpy
  import torch
  from omegaconf.dictconfig import DictConfig
  from nemo.collections.nlp.data.language_modeling.megatron.base_dataset_utils import get_datasets_weights_and_num_samples, get_train_valid_test_split_
  from nemo.collections.nlp.data.language_modeling.megatron.blendable_dataset import BlendableDataset
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_dataset import deallocate_indexed_dataset_memory
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_dataset import make_dataset
  from nemo.core import Dataset
  from nemo.utils import logging
label_desc: Use image node_id 6 for calling a global function with example usage: _build_index_mappings(name, data_prefix, documents, sizes, num_samples, seq_length, seed, index_mapping_dir, drop_last, add_extra_token, shuffle_documents, exchange_indices_distributed) and returns: doc_idx,  sample_idx,  shuffle_idx