code: def query_embeddings(
      db_dataset,
      index,
      embeddings,
      chunk_id_range,
      sample_map,
      n_chunks_per_sample,
      verbose=True,
  ):
      """Query neighbors of a block of embeddings."""
  
      args = get_retro_args()
  
      # Query neighbor ids.
      if verbose:
          print_rank_0("search.")
      t = time.time()
      assert (
          index.ntotal > 0
      ), "check we don't accidentally have an empty index."
      _, query_neighbor_ids = index.search(
          embeddings, args.retro_query_num_neighbors_query
      )
      if verbose:
          print_rank_0("  time : %.3f sec." % (time.time() - t))
  
      # Filter banned neighbor ids.
      if verbose:
          print_rank_0("filter banned neighbor ids.")
      filtered_neighbor_ids = np.full(
          shape=(
              len(query_neighbor_ids),
              args.retro_query_num_neighbors_save,
          ),
          fill_value=-1,
          dtype="int64",
      )
      min_chunk_id, max_chunk_id = chunk_id_range
      for chunk_id in range(min_chunk_id, max_chunk_id):
          sample_id = chunk_id // n_chunks_per_sample
          sample = sample_map[sample_id]
          sample_dataset_idx = sample["dataset_idx"].item()
          sample_doc_ids = sample["doc_ids"].tolist()
          sample_doc_tuples = [
              (sample_dataset_idx, d) for d in sample_doc_ids
          ]
  
          # Get valid neighbors (!= -1).
          query_row = [
              i
              for i in query_neighbor_ids[chunk_id - min_chunk_id]
              if i >= 0
          ]
  
          # Filter row.
          filtered_row = [
              i
              for i in query_row
              if tuple(db_dataset.doc_tuples[i].tolist())
              not in sample_doc_tuples
          ]
          filtered_row = filtered_row[
              : args.retro_query_num_neighbors_save
          ]
          filtered_row += [-1] * (
              args.retro_query_num_neighbors_save - len(filtered_row)
          )
          filtered_neighbor_ids[chunk_id - min_chunk_id] = filtered_row
  
      return query_neighbor_ids, filtered_neighbor_ids
  
imports: import numpy
  import os
  import psutil
  import time
  import torch
  from tqdm import tqdm
  from megatron import get_retro_args, print_rank_0
  from tools.bert_embedding import BertEmbedder
  from tools.bert_embedding.utils import get_missing_blocks_by_rank
  from tools.retro.db.utils import get_merged_train_dataset
  from tools.retro.external_libs import faiss, h5py
  from tools.retro.index.factory import IndexFactory
  from tools.retro.index.utils import get_index_dir
  from tools.retro.utils import GPTToTextDataset
  from .chunk_dataset import get_chunk_dataset_map
label_desc: Use image node_id 3 for calling a global function with example usage: query_embeddings(db_dataset, index, embeddings, chunk_id_range, sample_map, n_chunks_per_sample, verbose) and returns: query_neighbor_ids,  filtered_neighbor_ids