def _build_document_index(
    documents: numpy.ndarray,
    num_epochs: int,
    numpy_random_state: numpy.random.RandomState,
    separate_final_epoch: bool,
    enable_shuffle: bool = False,
) -> 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

        enable_shuffle (bool): Whether to enable the shuffle. Default is False to ensure the reproducibility

    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)
        if enable_shuffle:
            print("INFO: document_index shuffle is enabled...")
            numpy_random_state.shuffle(document_index)
        else:
            print("INFO: document_index shuffle is disabled...")
        return document_index

    doc_idx_first = _build_document_index(documents, num_epochs - 1, numpy_random_state, False, enable_shuffle)
    doc_idx_last = _build_document_index(documents, 1, numpy_random_state, False, enable_shuffle)
    return numpy.concatenate((doc_idx_first, doc_idx_last))
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

        enable_shuffle (bool): Whether to enable the shuffle. Default is False to ensure the reproducibility

    Returns:
        numpy.ndarray: The document index

    TODO: Explain separate_final_epoch
    
finetune/core/datasets/gpt_dataset.py