class_docstr: Class to build and search annoy indices.
  
  TODO(b/162421415): Split the managing of indices and the index logic.
  
  Upon instantiation of this class, it will first create empty indices for
  every model/embedding/dataset triple. It will then search for existing indices
  in the data directory and fill those. If the flag `initialize_new_indices` is
  set, it will fill the remaining indices by iterating over the data.
  During the saving process, both the index and a mapping from index-row to
  example are saved. These are used during the nearest neighbor lookup to return
  the closest example.
  
  Attributes:
     models: specification akin to the LIT server.
     datasets: specification akin to the LIT server.
     data_dir: path for (de-)serialization of indices.
     initialize_new_indices: whether to build new indices or simply load
       existing ones.
docstr: Find the nearest neighbor in index for an embedding.
  
  This function implements the search API for this class.
  The model/dataset/embedding combination maps to a unique index. Within, we
  are looking up the num_neighbors nearest neighbors of the embedding arg.
  
  Args:
    model_name: The identifier of the model.
    dataset_name: The identifier of the dataset.
    embedding_name: The identifier of the embedding within the model spec.
    embedding: The embedding we aim to look up.
    num_neighbors: How many neighbors we should return.
  
  Returns:
    neighbors: A list[dict] with nearest neighbor examples.
docstr_tok: ["Find","the","nearest","neighbor","in","index","for","an","embedding",".","This","function","implements","the","search","API","for","this","class",".","The","model\/dataset\/embedding","combination","maps","to","a","unique","index",".","Within",",","we","are","looking","up","the","num_neighbors","nearest","neighbors","of","the","embedding","arg",".","Args",":","model_name",":","The","identifier","of","the","model",".","dataset_name",":","The","identifier","of","the","dataset",".","embedding_name",":","The","identifier","of","the","embedding","within","the","model","spec",".","embedding",":","The","embedding","we","aim","to","look","up",".","num_neighbors",":","How","many","neighbors","we","should","return",".","Returns",":","neighbors",":","A","list","[","dict","]","with","nearest","neighbor","examples","."]
code: def find_nn(
      self,
      model_name: str,
      dataset_name: str,
      embedding_name: str,
      embedding: list[float],
      num_neighbors: Optional[int] = 25,
  ):
      """Find the nearest neighbor in index for an embedding.
  
      This function implements the search API for this class.
      The model/dataset/embedding combination maps to a unique index. Within, we
      are looking up the num_neighbors nearest neighbors of the embedding arg.
  
      Args:
        model_name: The identifier of the model.
        dataset_name: The identifier of the dataset.
        embedding_name: The identifier of the embedding within the model spec.
        embedding: The embedding we aim to look up.
        num_neighbors: How many neighbors we should return.
  
      Returns:
        neighbors: A list[dict] with nearest neighbor examples.
      """
      index_key = self._get_index_key(
          model_name, dataset_name, embedding_name
      )
      index = self._indices.get(index_key)
      assert index is not None, f"No index found for {index_key}."
  
      # Query for the neighbors.
      neighbor_indices, distances = index.get_nns_by_vector(
          vector=embedding, n=num_neighbors, include_distances=True
      )
  
      # Convert neighbors to texts.
      lookup_key = self._get_lookup_key(model_name, dataset_name)
      neighbor_examples = [
          self._example_lookup[lookup_key][ix]
          for ix in neighbor_indices
      ]
  
      # TODO(lit-dev): make the distance metadata that can be returned.
      del distances
  
      return neighbor_examples
  
imports: import collections
  from collections.abc import Mapping
  import os
  import pickle
  from typing import Optional
  from absl import logging
  import annoy
  from lit_nlp.api import dataset
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import utils
label_desc: Use image node_id 12 for calling the Indexer obj's underlying member method code with example usage: obj.find_nn(model_name, dataset_name, embedding_name, embedding, num_neighbors) and returns: neighbor_examples