code: def _build_doc_idx(
      documents, num_epochs, np_rng, separate_last_epoch, shuffle=True
  ):
      """Build an array with length = number-of-epochs * number-of-dcuments.
      Each index is mapped to a corresponding document."""
      if not separate_last_epoch or num_epochs == 1:
          doc_idx = np.mgrid[0:num_epochs, 0 : len(documents)][1]
          doc_idx[:] = documents
          doc_idx = doc_idx.reshape(-1)
          doc_idx = doc_idx.astype(np.int32)
          if shuffle:
              np_rng.shuffle(doc_idx)
          else:
              logging.info("Document shuffling disabled")
          return doc_idx
  
      doc_idx_first = _build_doc_idx(
          documents, num_epochs - 1, np_rng, False, shuffle
      )
      doc_idx_last = _build_doc_idx(
          documents, 1, np_rng, False, shuffle
      )
      return np.concatenate((doc_idx_first, doc_idx_last))
  
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 9 for calling a global function with example usage: _build_doc_idx(documents, num_epochs, np_rng, separate_last_epoch, shuffle) and returns: np, doc_idx