code: def build_individual_db(
      dataset_idx, n_datasets, dataset_info, tokenizers
  ):
      """Process a single indexed dataset & extract chunks."""
  
      args = get_retro_args()
  
      # Make directory.
      db_dir = dataset_info["db_dir"]
      os.makedirs(db_dir, exist_ok=True)
  
      # Indexed dataset.
      indexed_dataset = dataset_info["dataset"]
  
      # Missing db blocks.
      n_missing_world, missing_db_blocks = get_missing_blocks_by_rank(
          db_dir,
          len(indexed_dataset),
          args.retro_doc_block_size,
          validate=lambda f: f["chunks_valid"].shape == (0,)
          or f["chunks_valid"].shape[1] == 4,
      )
  
      # Prevent missing-path-write race condition.
      torch.distributed.barrier()
  
      if not missing_db_blocks:
          return
  
      # Num processes.
      if n_missing_world == 1:
          n_procs = 128
      elif n_missing_world <= 2:
          n_procs = 64
      elif n_missing_world <= 4:
          n_procs = 32
      elif n_missing_world <= 8:
          n_procs = 16
      else:
          n_procs = 8
  
      # Process documents in parallel.
      with ProcessPoolExecutor(max_workers=n_procs) as executor:
          for block_idx, block in enumerate(missing_db_blocks):
              if block is not None:
                  db_path = block["path"]
  
                  # Build partial dbs.
                  print_rank_0(" > build partial dbs.")
                  futures = []
                  for proc_id in range(n_procs):  # not true process id
                      futures.append(
                          executor.submit(
                              build_partial_db,
                              dataset_idx,
                              n_datasets,
                              indexed_dataset,
                              block_idx,
                              len(missing_db_blocks),
                              block,
                              proc_id,
                              n_procs,
                              tokenizers,
                          )
                      )
                  partial_chunk_dbs = []
                  for future in as_completed(futures):
                      partial_chunk_dbs.append(future.result())
  
                  # Concatenate chunks.
                  partial_chunk_dbs.sort(
                      key=lambda item: item[0]
                  )  # sort by proc_id
                  chunk_db_valid = [
                      item
                      for partial_chunk_db in partial_chunk_dbs
                      for item in partial_chunk_db[1]
                  ]
                  chunk_db_invalid = [
                      item
                      for partial_chunk_db in partial_chunk_dbs
                      for item in partial_chunk_db[2]
                  ]
  
                  # Convert to numpy.
                  print_rank_0(" > converting chunk db to numpy.")
                  chunk_db_valid = np.array(
                      chunk_db_valid, dtype="uint32"
                  )
                  chunk_db_invalid = np.array(
                      chunk_db_invalid, dtype="uint32"
                  )
  
                  # Document offsets.
                  doc_sizes = [
                      (d, s)
                      for partial_chunk_db in partial_chunk_dbs
                      for d, s in partial_chunk_db[3].items()
                  ]
                  doc_sizes.sort(key=lambda item: item[0])
                  doc_offsets = np.cumsum(
                      [item[1] for item in doc_sizes]
                  ).astype("uint64")
                  doc_offsets = np.stack(
                      (
                          np.array(
                              [item[0] for item in doc_sizes],
                              dtype="uint64",
                          ),
                          doc_offsets,
                      ),
                      axis=1,
                  )
  
                  # Save DB.
                  print_rank_0(" > saving individual db.")
                  with h5py.File(db_path, "w") as f:
                      dset = f.create_dataset(
                          "chunks_valid", data=chunk_db_valid
                      )
                      dset = f.create_dataset(
                          "chunks_invalid", data=chunk_db_invalid
                      )
                      dset = f.create_dataset(
                          "doc_offsets", data=doc_offsets
                      )
  
              # Wait for all ranks to finish block.
              print_rank_0(" > waiting for all ranks to finish block.")
              torch.distributed.barrier()
  
      print_rank_0(" > finished saving individual db.")
  
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 3 for calling a global function with example usage: build_individual_db(dataset_idx, n_datasets, dataset_info, tokenizers) without return types