class_docstr: Computes nearest neighbors of an example embedding.
  
  Required Model Output:
    - Embeddings (`emb_layer`) to return the input embeddings
        for a layer
docstr: Finds the nearest neighbors of the example specified in the config.
  
  Args:
    inputs: the dataset example to find nearest neighbors for.
    model: the model being explained.
    dataset: the dataset which the current examples belong to.
    model_outputs: optional model outputs from calling model.predict(inputs).
    config: a config which should specify:
      {
        'num_neighbors': [the number of nearest neighbors to return]
        'embedding_name': [the name of the embedding field to use]
        'use_input': [Optional boolean if the embedding comes from input data]
      }
  
  Raises:
    KeyError: Cannot find the embedding field in the relevant spec
    TypeError: `config` argument not provided
    ValueError: indexed_inputs requires exactly one input
  
  Returns:
    A JsonDict containing the a list of num_neighbors nearest neighbors,
    where each has the example id and distance from the main example.
docstr_tok: ["Finds","the","nearest","neighbors","of","the","example","specified","in","the","config",".","Args",":","inputs",":","the","dataset","example","to","find","nearest","neighbors","for",".","model",":","the","model","being","explained",".","dataset",":","the","dataset","which","the","current","examples","belong","to",".","model_outputs",":","optional","model","outputs","from","calling","model.predict","(","inputs",")",".","config",":","a","config","which","should","specify",":","{","'num_neighbors","'",":","[","the","number","of","nearest","neighbors","to","return","]","'embedding_name","'",":","[","the","name","of","the","embedding","field","to","use","]","'use_input","'",":","[","Optional","boolean","if","the","embedding","comes","from","input","data","]","}","Raises",":","KeyError",":","Can","not","find","the","embedding","field","in","the","relevant","spec","TypeError",":","`","config","`","argument","not","provided","ValueError",":","indexed_inputs","requires","exactly","one","input","Returns",":","A","JsonDict","containing","the","a","list","of","num_neighbors","nearest","neighbors",",","where","each","has","the","example","id","and","distance","from","the","main","example","."]
code: def run(
      self,
      inputs: Sequence[JsonDict],
      model: lit_model.Model,
      dataset: lit_dataset.Dataset,
      model_outputs: Optional[list[JsonDict]] = None,
      config: Optional[JsonDict] = None,
  ) -> Optional[list[JsonDict]]:
      """Finds the nearest neighbors of the example specified in the config.
  
      Args:
        inputs: the dataset example to find nearest neighbors for.
        model: the model being explained.
        dataset: the dataset which the current examples belong to.
        model_outputs: optional model outputs from calling model.predict(inputs).
        config: a config which should specify:
          {
            'num_neighbors': [the number of nearest neighbors to return]
            'embedding_name': [the name of the embedding field to use]
            'use_input': [Optional boolean if the embedding comes from input data]
          }
  
      Raises:
        KeyError: Cannot find the embedding field in the relevant spec
        TypeError: `config` argument not provided
        ValueError: indexed_inputs requires exactly one input
  
      Returns:
        A JsonDict containing the a list of num_neighbors nearest neighbors,
        where each has the example id and distance from the main example.
      """
      if not config:
          raise TypeError("config must be provided")
  
      if not (isinstance(dataset, lit_dataset.IndexedDataset)):
          raise TypeError(
              "Nearest neighbors requires an IndexedDataset to track "
              "uniqueness by ID."
          )
  
      nnconf = NearestNeighborsConfig(
          **{k: v for k, v in config.items() if k in _NN_CONFIG_FIELDS}
      )
  
      # TODO(lit-dev): Add support for selecting nearest neighbors of a set.
      if len(inputs) != 1:
          raise ValueError(
              "indexed_inputs must contain exactly 1 example, found "
              f"{len(inputs)}."
          )
  
      if nnconf.use_input:
          if not dataset.spec().get(nnconf.embedding_name):
              raise KeyError(
                  "Could not find embeddings field, "
                  f"{nnconf.embedding_name} in dataset spec"
              )
          # If using input values, then treat inputs as outputs instead of running
          # the model.
          dataset_outputs = dataset.examples
          example_outputs = inputs
      else:
          if not model.output_spec().get(nnconf.embedding_name):
              raise KeyError(
                  "Could not find embeddings field, "
                  f"{nnconf.embedding_name} in model output spec"
              )
          dataset_outputs = list(model.predict(dataset.examples))
          example_outputs = list(model.predict(inputs))
  
      example_output = example_outputs[0]
  
      # <float32>[emb_size]
      dataset_embs = [
          output[nnconf.embedding_name] for output in dataset_outputs
      ]
      example_embs = [example_output[nnconf.embedding_name]]
      distances = distance.cdist(example_embs, dataset_embs)[0]
      sorted_indices = np.argsort(distances)
      k = nnconf.num_neighbors
      k_nearest_neighbors = [
          {
              "id": dataset.examples[original_index]["_id"],
              "nn_distance": distances[original_index],
          }
          for original_index in sorted_indices[:k]
      ]
  
      return [{"nearest_neighbors": k_nearest_neighbors}]
  
file_path: lit/lit_nlp/components/nearest_neighbors.py
imports: from collections.abc import Sequence
  import dataclasses
  from typing import Optional
  from lit_nlp.api import components
  from lit_nlp.api import dataset
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import utils
  import numpy
  from scipy.spatial import distance
label_desc: Use image node_id 2 for calling the NearestNeighbors obj's underlying member method code with example usage: obj.run(inputs, model, dataset, model_outputs, config) and returns: list