class_docstr: ## Index for retrieving nearest neighbors
class_docstr_tok: ["#","#","Index","for","retrieving","nearest","neighbors"]
docstr: * `chunk_len` is the chunk length
  * `n_probe` is the number of lists to probe
  * `n_neighbors` is the number of neighbors to retrieve
  * `n_extra` is the number of extra neighbors to retrieve since we will be
      removing neighbors overlapping with the query chunk
  * `exclude_neighbor_span` is the extra text length to avoid when checking for overlaps
docstr_tok: ["*","`","chunk_len","`","is","the","chunk","length","*","`","n_probe","`","is","the","number","of","lists","to","probe","*","`","n_neighbors","`","is","the","number","of","neighbors","to","retrieve","*","`","n_extra","`","is","the","number","of","extra","neighbors","to","retrieve","since","we","will","be","removing","neighbors","overlapping","with","the","query","chunk","*","`","exclude_neighbor_span","`","is","the","extra","text","length","to","avoid","when","checking","for","overlaps"]
code: def __init__(
      self,
      chunk_len: int = 16,
      n_probe: int = 8,
      n_neighbors: int = 2,
      n_extra: int = 2,
      exclude_neighbor_span: int = 8,
  ):
      """
      * `chunk_len` is the chunk length
      * `n_probe` is the number of lists to probe
      * `n_neighbors` is the number of neighbors to retrieve
      * `n_extra` is the number of extra neighbors to retrieve since we will be
          removing neighbors overlapping with the query chunk
      * `exclude_neighbor_span` is the extra text length to avoid when checking for overlaps
      """
  
      self.n_neighbors = n_neighbors
      self.chunk_len = chunk_len
      self.exclude_neighbor_span = exclude_neighbor_span
      self.n_extra = n_extra
  
      # Initialize BERT to get $\text{B\small{ERT}}(N)$
      self.bert = BERTChunkEmbeddings(torch.device("cuda:0"))
      # Load the database
      with monit.section("Load index"):
          self.index = faiss.read_index(
              str(lab.get_data_path() / "retro.index")
          )
          self.index.nprobe = n_probe
  
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 1 to create a new RetroIndex object  with example: obj = RetroIndex(chunk_len, n_probe, n_neighbors, n_extra, exclude_neighbor_span)