class_docstr: ## Index for retrieving nearest neighbors
class_docstr_tok: ["#","#","Index","for","retrieving","nearest","neighbors"]
docstr: ### Retrieve nearest neighbors
docstr_tok: ["#","#","#","Retrieve","nearest","neighbors"]
code: def __call__(
      self, query_chunks: List[str], offsets: Optional[List[int]]
  ):
      """
      ### Retrieve nearest neighbors
      """
  
      # Get $\text{B\small{ERT}}(N)$ of query chunks
      emb = self.bert(query_chunks).cpu()
  
      # Get `n_neighbors + n_extra` nearest neighbors from the database
      distance, neighbor_offsets = self.index.search(
          emb.numpy(), self.n_neighbors + self.n_extra
      )
  
      # If the query chunk offsets are given filter out overlapping chunks
      if offsets is not None:
          neighbor_offsets = [
              self.filter_neighbors(off, n_off)
              for off, n_off in zip(offsets, neighbor_offsets)
          ]
  
      # Get the closest `n_neighbors` after filtering
      neighbor_offsets = [
          n_off[: self.n_neighbors] for n_off in neighbor_offsets
      ]
  
      #
      return neighbor_offsets
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/retro/database.py
imports: from typing import List, Optional
  import faiss
  import numpy
  import torch
  from labml import lab, monit
  from labml_helpers.datasets.text import TextFileDataset
  from labml_nn.transformers.retro.bert_embeddings import BERTChunkEmbeddings
label_desc: Use image node_id 3 for calling the RetroIndex obj's underlying member method code with example usage: obj.__call__(query_chunks, offsets) and returns: neighbor_offsets