code: def load_datapool(dpath):
      def load_single_file(saved_embeddings):
          compressed = np.load(saved_embeddings)
          database = {key: compressed[key] for key in compressed.files}
          return database
  
      def load_multi_files(data_archive):
          database = {key: [] for key in data_archive[0].files}
          for d in tqdm(
              data_archive,
              desc=f"Loading datapool from {len(data_archive)} individual files.",
          ):
              for key in d.files:
                  database[key].append(d[key])
  
          return database
  
      print(f'Load saved patch embedding from "{dpath}"')
      file_content = glob.glob(os.path.join(dpath, "*.npz"))
  
      if len(file_content) == 1:
          data_pool = load_single_file(file_content[0])
      elif len(file_content) > 1:
          data = [np.load(f) for f in file_content]
          prefetched_data = parallel_data_prefetch(
              load_multi_files,
              data,
              n_proc=min(len(data), cpu_count()),
              target_data_type="dict",
          )
  
          data_pool = {
              key: np.concatenate(
                  [od[key] for od in prefetched_data], axis=1
              )[0]
              for key in prefetched_data[0].keys()
          }
      else:
          raise ValueError(
              f'No npz-files in specified path "{dpath}" is this directory existing?'
          )
  
      print(
          f'Finished loading of retrieval database of length {data_pool["embedding"].shape[0]}.'
      )
      return data_pool
  
file_path: ColossalAI/examples/images/diffusion/scripts/train_searcher.py
imports: import argparse
  import glob
  import os
  import sys
  from multiprocessing import cpu_count
  import numpy
  import scann
  from ldm.util import parallel_data_prefetch
  from tqdm import tqdm
label_desc: Use image node_id 4 for calling a global function with example usage: load_datapool(dpath) and returns: data_pool, database, database