Nearest Neighbor Normalization documentation

Here, you will find documentation for every API method in the nnn-retrieval pip package. See https://github.com/multimodal-interpretability/nnn for an overview of the package.

Contents

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).

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).

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).

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

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

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

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