class_docstr: Wrapper object for a BlockData which similarity search via FAISS under the hood
class_docstr_tok: ["Wrapper","object","for","a","BlockData","which","similarity","search","via","FAISS","under","the","hood"]
docstr: Add the embedding of each block to the underlying FAISS index
docstr_tok: ["Add","the","embedding","of","each","block","to","the","underlying","FAISS","index"]
code: def add_embed_data(self, all_embed_data):
      """Add the embedding of each block to the underlying FAISS index"""
  
      # this assumes the embed_data is a dict : {int: np.array<float>}
      block_indices, block_embeds = zip(
          *all_embed_data.embed_data.items()
      )
  
      # the embeddings have to be entered in as float32 even though the math
      # internally is done with float16.
      embeds_arr = np.float32(np.array(block_embeds))
      indices_arr = np.array(block_indices)
  
      # we no longer need the embedding data since it's in the index now
      all_embed_data.clear()
  
      self.mips_index.add_with_ids(embeds_arr, indices_arr)
  
      if (
          not mpu.model_parallel_is_initialized()
          or mpu.get_data_parallel_rank() == 0
      ):
          print(">>> Finished adding block data to index", flush=True)
  
imports: import itertools
  import os
  import pickle
  import shutil
  import numpy
  import torch
  from megatron import get_args
  from megatron.core import mpu
label_desc: Use image node_id 5 for calling the FaissMIPSIndex obj's underlying member method code with example usage: obj.add_embed_data(all_embed_data) without return types