def __init__(self, query_embedder, documents, doc_embeds_or_path,
               num_neighbors):
    """Constructs BruteForceRetriever.

    Args:
      query_embedder: an instance of QueryEmbedder.
      documents: a list of Document objects.
      doc_embeds_or_path: either a [num_docs, embed_dim] TF Tensor, or a path to
        load it.
      num_neighbors: number of neighbors to retrieve.
    """
    total_candidates = len(documents)

    self._query_embedder = query_embedder
    self._num_neighbors = num_neighbors
    self._documents = documents

    # Load embeddings.
    if isinstance(doc_embeds_or_path, str):
      with tf.device('/CPU:0'):
        ckpt_reader = tf.train.load_checkpoint(doc_embeds_or_path)
        self._doc_embeds = ckpt_reader.get_tensor('block_emb')
    else:
      self._doc_embeds = doc_embeds_or_path
    logging.info('Loaded document embeddings.')

    # Check shapes.
    if self._doc_embeds.shape[0] != total_candidates:
      raise ValueError('Did not load the right number of embeddings.')
Constructs BruteForceRetriever.

    Args:
      query_embedder: an instance of QueryEmbedder.
      documents: a list of Document objects.
      doc_embeds_or_path: either a [num_docs, embed_dim] TF Tensor, or a path to
        load it.
      num_neighbors: number of neighbors to retrieve.
    
language/realm/retrieval.py