nnn package

Submodules

nnn.base_ranker module

class nnn.base_ranker.BaseRanker(retriever, retrieval_embeds, reference_embeds, batch_size=128, use_gpu=False, gpu_id=-1)

Bases: Ranker

__init__(retriever, retrieval_embeds, reference_embeds, batch_size=128, use_gpu=False, gpu_id=-1)
search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

nnn.base_retriever module

class nnn.base_retriever.BaseRetriever(embeds_size, use_gpu=False, gpu_id=-1)

Bases: Retriever

__init__(embeds_size, use_gpu=False, gpu_id=-1)

Initializes a new instance of BaseRetriever.

Parameters:

  • embeds_size (int) – The dimension size of embeddings.

  • use_gpu (bool, optional) – If True, will attempt to use GPU. Default is False.

  • gpu_id (int, optional) – Specifies GPU device ID if use_gpu is True. Default is -1 (no GPU device).

Raises:

Exception – If use_gpu is True but gpu_id is not specified.

compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

nnn.dn_ranker module

class nnn.dn_ranker.DNRanker(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1, retrieval_dev_embeds=None, query_dev_embeds=None, dn_lambda=0.5)

Bases: Ranker

__init__(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1, retrieval_dev_embeds=None, query_dev_embeds=None, dn_lambda=0.5)
search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries, using the test-time distribution normalization method to modify the similarity scores.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

nnn.faiss_cpu_retriever module

class nnn.faiss_cpu_retriever.FaissCPURetriever(embeds_size, reference_index=None, reference_nprobes=32, retrieval_index=None, retrieval_nprobes=32)

Bases: Retriever

__init__(embeds_size, reference_index=None, reference_nprobes=32, retrieval_index=None, retrieval_nprobes=32)

Initializes a new instance of FaissCPURetriever.

Parameters:

  • embeds_size (int) – The dimension size of embeddings.

  • reference_index (Faiss index) – a Faiss index to be used to search reference embeddings. Should use inner product metric, as this is the distance metric we are operating in.

  • reference_nprobes (int) – See nprobes for Faiss indices.

  • retrieval_index (Faiss index) – a Faiss index to be used to search reference embeddings. Should use inner product metric, as this is the distance metric we are operating in.

  • retrieval_nprobes (int) – See nprobes for Faiss indices.

check_dimensions(retrieval_embeds)
compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

nnn.faiss_gpu_retriever module

class nnn.faiss_gpu_retriever.FaissGPURetriever(embeds_size, gpu_id, reference_index=None, reference_nprobes=32, retrieval_index=None, retrieval_nprobes=32)

Bases: Retriever

__init__(embeds_size, gpu_id, reference_index=None, reference_nprobes=32, retrieval_index=None, retrieval_nprobes=32)

Initializes a new instance of FaissCPURetriever.

Parameters:

  • embeds_size (int) – The dimension size of embeddings.

  • gpu_id (int) – Specifies GPU device ID if use_gpu is True. Default is -1 (no GPU device).

  • reference_index (Faiss index) – a Faiss index to be used to search reference embeddings. Should use inner product metric, as this is the distance metric we are operating in.

  • reference_nprobes (int) – See nprobes for Faiss indices.

  • retrieval_index (Faiss index) – a Faiss index to be used to search reference embeddings. Should use inner product metric, as this is the distance metric we are operating in.

  • retrieval_nprobes (int) – See nprobes for Faiss indices.

compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

nnn.nnn_ranker module

class nnn.nnn_ranker.NNNRanker(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1)

Bases: Ranker

__init__(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1)
search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries, using the test-time nearest neighbor normalization method to modify the similarity scores.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

nnn.nnn_retriever module

class nnn.nnn_retriever.NNNRetriever(embeds_size, use_gpu=False, gpu_id=-1)

Bases: Retriever

__init__(embeds_size, use_gpu=False, gpu_id=-1)

Initializes a new instance of NNNRetriever.

Parameters:

  • embeds_size (int) – The dimension size of embeddings.

  • use_gpu (bool, optional) – If True, will attempt to use GPU. Default is False.

  • gpu_id (int, optional) – Specifies GPU device ID if use_gpu is True. Default is -1 (no GPU device).

Raises:

Exception – If use_gpu is True but gpu_id is not specified.

compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

nnn.ranker module

class nnn.ranker.Ranker

Bases: ABC

In our framework, the Ranker class is responsible for the defining the method of score corrections to be used during retrieval. Currently supported methods of score corrections are our proposed method of Nearest Neighbor Normalization, Distribution Normalization (proposed in https://arxiv.org/abs/2302.11084), and a base no-correction method.

A ranker object is initialized with a retriever object, and orchestrates the actual search process by computing the necessary score corrections, then using the retriever object to perform search with these score corrections. See https://github.com/multimodal-interpretability/nnn for examples of use.

abstract search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Return type:

Tuple[matrix, matrix]

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

nnn.retriever module

class nnn.retriever.Retriever

Bases: ABC

In our framework, the Retriever class is responsible for the actual retrieval mechanisms. It is intended to define the method of retrieval (exhaustive search using Pytorch, narest-neighbor vector search using Faiss), and should be used with a Ranker class that defines the actual method of score corrections to be used during retrieval (which includes Nearest Neighbor Normalization). See the docs for the Ranker class for more detail.

Note that the retriever class does not store any information locally, and depends on passed in values for retrieval embeddings, query embeddings, reference embeddings, and/or additive corrections, which what define the actual Nearest Neighbor Normalization method. Thus, it should be used with a ranker object, rather than directly used.

abstract compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

abstract retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

abstract setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

Module contents

class nnn.BaseRanker(retriever, retrieval_embeds, reference_embeds, batch_size=128, use_gpu=False, gpu_id=-1)

Bases: Ranker

__init__(retriever, retrieval_embeds, reference_embeds, batch_size=128, use_gpu=False, gpu_id=-1)
search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

class nnn.BaseRetriever(embeds_size, use_gpu=False, gpu_id=-1)

Bases: Retriever

__init__(embeds_size, use_gpu=False, gpu_id=-1)

Initializes a new instance of BaseRetriever.

Parameters:

  • embeds_size (int) – The dimension size of embeddings.

  • use_gpu (bool, optional) – If True, will attempt to use GPU. Default is False.

  • gpu_id (int, optional) – Specifies GPU device ID if use_gpu is True. Default is -1 (no GPU device).

Raises:

Exception – If use_gpu is True but gpu_id is not specified.

compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

class nnn.DNRanker(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1, retrieval_dev_embeds=None, query_dev_embeds=None, dn_lambda=0.5)

Bases: Ranker

__init__(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1, retrieval_dev_embeds=None, query_dev_embeds=None, dn_lambda=0.5)
search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries, using the test-time distribution normalization method to modify the similarity scores.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

class nnn.NNNRanker(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1)

Bases: Ranker

__init__(retriever, retrieval_embeds, reference_embeds, alternate_ks=256, batch_size=128, alternate_weight=0.5, use_gpu=False, gpu_id=-1)
search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries, using the test-time nearest neighbor normalization method to modify the similarity scores.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

class nnn.NNNRetriever(embeds_size, use_gpu=False, gpu_id=-1)

Bases: Retriever

__init__(embeds_size, use_gpu=False, gpu_id=-1)

Initializes a new instance of NNNRetriever.

Parameters:

  • embeds_size (int) – The dimension size of embeddings.

  • use_gpu (bool, optional) – If True, will attempt to use GPU. Default is False.

  • gpu_id (int, optional) – Specifies GPU device ID if use_gpu is True. Default is -1 (no GPU device).

Raises:

Exception – If use_gpu is True but gpu_id is not specified.

compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

class nnn.Ranker

Bases: ABC

In our framework, the Ranker class is responsible for the defining the method of score corrections to be used during retrieval. Currently supported methods of score corrections are our proposed method of Nearest Neighbor Normalization, Distribution Normalization (proposed in https://arxiv.org/abs/2302.11084), and a base no-correction method.

A ranker object is initialized with a retriever object, and orchestrates the actual search process by computing the necessary score corrections, then using the retriever object to perform search with these score corrections. See https://github.com/multimodal-interpretability/nnn for examples of use.

abstract search(batch_query, top_k)

Searches for the top_k most similar items to the given batch of queries.

Parameters:

  • batch_query (np.matrix) – A batch of query embeddings. Dimensions are (n_queries, embedding_dim).

  • top_k (int) – The number of top similar items to retrieve.

Return type:

Tuple[matrix, matrix]

Returns:

The method should return an array of indices of the top_k most similar items in order for each query in the batch, with dimensions (n_queries, top_k).

class nnn.Retriever

Bases: ABC

In our framework, the Retriever class is responsible for the actual retrieval mechanisms. It is intended to define the method of retrieval (exhaustive search using Pytorch, narest-neighbor vector search using Faiss), and should be used with a Ranker class that defines the actual method of score corrections to be used during retrieval (which includes Nearest Neighbor Normalization). See the docs for the Ranker class for more detail.

Note that the retriever class does not store any information locally, and depends on passed in values for retrieval embeddings, query embeddings, reference embeddings, and/or additive corrections, which what define the actual Nearest Neighbor Normalization method. Thus, it should be used with a ranker object, rather than directly used.

abstract compute_alignment_means(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Computes mean similarity scores between per-retrieval embedding, using a given set reference embeddings. As described in the paper, these mean scores are later used as an additive correction to the base retrieval scores.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average across.i

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray

abstract retrieve(retrieval_embeds, batch_query, top_k, alternate_weight, alignment_means, batch_size)

Not intended to be a public facing method; all retrieval should be done from the Ranker interface.

Retrieves the top_k most similar items for a batch of query embeddings.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • batch_query (torch.Tensor) – Query embeddings to retrieve items for.

  • top_k (int) – Number of top items to retrieve.

  • alternate_weight (float) – Weight to adjust alignment means in similarity scores.

  • alignment_means (torch.Tensor) – Precomputed alignment means per retrieval embedding to adjust retrieval scores.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Contains two numpy arrays:

  • distances: The top-k similarity scores for each query.

  • indices: The indices of the top-k retrieved items.

Return type:

tuple

abstract setup_retriever(retrieval_embeds, reference_embeds, alternate_ks, batch_size)

Not intended to be a public facing method; should be called internally in the Ranker interface.

Sets up the retriever by calculating alignment means for retrieval, allowing for alignment means to be used at inference time for retrieval as described in the paper.

Parameters:

  • retrieval_embeds (torch.Tensor) – Embeddings from the retrieval set.

  • reference_embeds (torch.Tensor) – Reference embeddings to compare against.

  • alternate_ks (int) – Number of top-k reference scores to average.

  • batch_size (int) – Number of samples to process per batch.

Returns:

Array of alignment means per retrieval embedding.

Return type:

numpy.ndarray