code: def _build_document_index(
      documents: numpy.ndarray,
      num_epochs: int,
      numpy_random_state: numpy.random.RandomState,
      separate_final_epoch: bool,
  ) -> numpy.ndarray:
      """Build an array with length = num epochs * num documents
  
      Args:
          documents (numpy.ndarray): the subset of exposed document indices
  
          num_epochs (int): The number of epochs
  
          numpy_random_state (numpy.random.RandomState): The NumPy random state
  
          separate_final_epoch (bool): Whether to exclude the last epoch from the global shuffle
  
      Returns:
          numpy.ndarray: The document index
  
      TODO: Explain separate_final_epoch
      """
      if not separate_final_epoch or num_epochs == 1:
          document_index = numpy.mgrid[
              0:num_epochs, 0 : len(documents)
          ][1]
          document_index[:] = documents
          document_index = document_index.reshape(-1)
          document_index = document_index.astype(numpy.int32)
          numpy_random_state.shuffle(document_index)
          return document_index
  
      doc_idx_first = _build_document_index(
          documents, num_epochs - 1, numpy_random_state, False
      )
      doc_idx_last = _build_document_index(
          documents, 1, numpy_random_state, False
      )
      return numpy.concatenate((doc_idx_first, doc_idx_last))
  
imports: import logging
  import os
  import time
  from dataclasses import dataclass
  from typing import Dict, Tuple, Union
  import numpy
  import torch
  from megatron.core.datasets.blended_megatron_dataset_config import BlendedMegatronDatasetConfig
  from megatron.core.datasets.indexed_dataset import MMapIndexedDataset
  from megatron.core.datasets.megatron_dataset import MegatronDataset
  from megatron.core.datasets.utils import Split, log_single_rank
label_desc: Use image node_id 1 for calling a global function with example usage: _build_document_index(documents, num_epochs, numpy_random_state, separate_final_epoch) and returns: numpy, document_index