code: def build_partial_db(
      dataset_idx,
      n_datasets,
      indexed_dataset,
      block_id,
      n_blocks,
      block,
      proc_id,
      n_procs,
      tokenizers,
  ):
      """Process a document index range of the indexed dataset.
  
      The chunk database is built in parallel blocks, since de-tokenizing &
      re-tokenizing for Bert-length computation is expensive. This method
      iterates each document and extracts sequential 'chunk-length' sequences
      from each document.
      """
  
      args = get_retro_args()
  
      # Document start/end indexes.
      doc_range = block["range"]
      n_docs = doc_range[1] - doc_range[0]
      n_docs_per_proc = int(np.ceil(n_docs / n_procs))
      doc_start_id = doc_range[0] + proc_id * n_docs_per_proc
      doc_end_id = min(doc_range[1], doc_start_id + n_docs_per_proc)
  
      # Print progress.
      progress_proc_ids = (
          set(range(n_procs))
          if torch.distributed.get_rank() == 0
          else set()
      )
      if proc_id in progress_proc_ids:
          print(
              " > building partial chunk db, proc %d / %d, docs %d:%d / %d."
              % (
                  proc_id,
                  n_procs,
                  doc_start_id,
                  doc_end_id,
                  n_docs,
              )
          )
  
      # Progress bars (snapshot of overall progress).
      doc_id_iter = range(doc_start_id, doc_end_id)
      pbar = (
          tqdm(doc_id_iter)
          if proc_id in progress_proc_ids
          else doc_id_iter
      )
  
      # Iterate documents & parse chunks.
      chunk_db_valid = []
      chunk_db_invalid = []
      doc_size_map = {}
      for doc_id in pbar:
          # Progress description.
          try:
              pbar.set_description(
                  "ds %d / %d, block %d / %d, proc %d / %d."
                  % (
                      dataset_idx,
                      n_datasets,
                      block_id,
                      n_blocks,
                      proc_id,
                      n_procs,
                  )
              )
          except:
              pass
  
          # Remove EOD token.
          doc = indexed_dataset.get(doc_id)
          if doc[-1].item() == tokenizers.gpt.eod:
              doc = doc[:-1]
          doc_len = len(doc)
  
          # Chunk start/end indexes.
          chunk_start_idxs = list(
              range(0, doc_len, args.retro_gpt_chunk_length)
          )
          chunk_end_idxs = [
              min(doc_len, s + args.retro_gpt_chunk_length)
              for s in chunk_start_idxs
          ]
  
          # Re-tokenize each chunk to Bert/Wordpiece (empty bert -> 'invalid').
          doc_size_map[doc_id] = 0
          for i, chunk_start_idx in enumerate(chunk_start_idxs):
              # Re-tokenize.
              chunk_end_idx = chunk_end_idxs[i]
              gpt_token_ids = indexed_dataset.get(
                  idx=doc_id,
                  offset=chunk_start_idx,
                  length=chunk_end_idx - chunk_start_idx,
              )
              text = tokenizers.gpt.detokenize(gpt_token_ids.tolist())
              bert_token_ids = tokenizers.bert.tokenize(text)
  
              # 'Valid' for non-empty Bert chunks; 'invalid' otherwise.
              if len(bert_token_ids) == 0:
                  _chunk_db = chunk_db_invalid
              else:
                  _chunk_db = chunk_db_valid
                  doc_size_map[doc_id] += 1
              _chunk_db.append(
                  (
                      doc_id,
                      chunk_start_idx,
                      chunk_end_idx,
                      len(bert_token_ids),
                  )
              )
  
      return proc_id, chunk_db_valid, chunk_db_invalid, doc_size_map
  
imports: from collections import defaultdict
  from concurrent.futures import as_completed, ProcessPoolExecutor
  from functools import reduce
  import glob
  import json
  import numpy
  import os
  from pathlib import Path
  import threading
  import torch
  from tqdm import tqdm
  import types
  from megatron import get_retro_args, print_rank_0
  from megatron.core.datasets.indexed_dataset import MMapIndexedDataset
  from megatron.tokenizer.tokenizer import _BertWordPieceTokenizer, _GPT2BPETokenizer
  from tools.bert_embedding.utils import get_missing_blocks_by_rank
  from tools.retro.external_libs import h5py
  from tools.retro.utils import get_gpt_tokenizer, get_bert_tokenizer
  from .utils import get_indexed_dataset_infos, get_indexed_dataset_infos_path, get_individual_db_dir, get_individual_chunk_db, get_individual_doc_offsets, get_merged_dataset, get_merged_db_path_map, save_indexed_dataset_infos
label_desc: Use image node_id 2 for calling a global function with example usage: build_partial_db(dataset_idx, n_datasets, indexed_dataset, block_id, n_blocks, block, proc_id, n_procs, tokenizers) and returns: proc_id,  chunk_db_valid,  chunk_db_invalid,  doc_size_map