code: def update_chunk_counts(indexed_dataset_infos):
      """Set n_chunks_train & n_chunks sampled for each individual DB."""
  
      args = get_retro_args()
  
      if torch.distributed.get_rank() != 0:
          return
  
      # Data ratio sum (for setting index training chunks).
      data_ratio_sum = sum([d["ratio"] for d in indexed_dataset_infos])
  
      # Training split size (split at document level).
      train_fraction = float(args.split.split(",")[0]) / 100
      assert train_fraction > 0 and train_fraction <= 1
  
      # Set n_chunks (including n_chunks_sampled for unambiguity).
      print_rank_0(" > compute n_chunks.")
      for ds_index, ds_info in enumerate(indexed_dataset_infos):
          db_dir = ds_info["db_dir"]
          db_paths = sorted(glob.glob(db_dir + "/*.hdf5"))
  
          # Update counts.
          ds_info["n_docs"] = (
              len(ds_info["dataset"].document_indices) - 1
          )
          ds_info["n_docs_train"] = int(
              train_fraction * ds_info["n_docs"]
          )
          ds_info["n_chunks"] = 0  # previously, 'n_chunks_valid'
          ds_info["n_chunks_train"] = 0
          ds_info["n_chunks_invalid"] = 0
          for db_path in tqdm(
              db_paths,
              "%d/%d, %s"
              % (ds_index, len(indexed_dataset_infos), ds_info["name"]),
          ):
              with h5py.File(db_path, "r") as f:
                  ds_info["n_chunks"] += len(f["chunks_valid"])
                  ds_info["n_chunks_invalid"] += len(
                      f["chunks_invalid"]
                  )
                  ds_info["n_chunks_train"] += (
                      (
                          np.copy(f["chunks_valid"][:, 0])
                          < ds_info["n_docs_train"]
                      )
                      .sum()
                      .item()
                  )
  
          ds_info["n_chunks_sampled"] = int(
              args.retro_index_ntrain
              * ds_info["ratio"]
              / data_ratio_sum
          )
  
          # Verify counts.
          assert (
              ds_info["n_chunks_train"] <= ds_info["n_chunks"]
          ), "n_train (%d) > n_total (%d)." % (
              ds_info["n_chunks_train"],
              ds_info["n_chunks"],
          )
          assert (
              ds_info["n_chunks_sampled"] <= ds_info["n_chunks_train"]
          ), "n_sampled (%d) > n_train (%d)." % (
              ds_info["n_chunks_sampled"],
              ds_info["n_chunks_train"],
          )
  
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 5 for calling a global function with example usage: update_chunk_counts(indexed_dataset_infos) without return types