Vendi Score

You can install the core package via pip:

pip install vendi_score

Alternatively, you can install it from source by cloning the repository:

git clone https://github.com/vertaix/Vendi-Score.git
cd Vendi-Score
pip install -e .

To use specialized similarity functions for images, text, or molecules, install the corresponding optional dependencies:

• Images: pip install vendi_score[images]
• Text and Molecules: pip install vendi_score[text,molecules]
• All extras: pip install vendi_score[all]

The library uses two primary data abstractions to manage diversity evaluations:

1. data_utils.Example: A utility class that wraps a single data point. The raw content (e.g., a sentence) is stored in the .x attribute, and metadata/labels are stored in a .labels dictionary.
2. data_utils.Group: A container for a collection of Example objects. It tracks:
   • .name: The name of the group.
   • .examples: A list of Example objects.
   • .features: A dictionary storing feature matrices (e.g., n-gram counts).
   • .Ks: A dictionary storing similarity matrices (e.g., $X \times X^T$).
   • .metrics: A dictionary storing calculated diversity scores (e.g., Vendi Score).

The data_utils module provides helpers for organizing data and simulating specific scenarios:

• data_utils.Group: A container to hold a set of data_utils.Example objects, their associated labels, and calculated metrics.
• data_utils.mode_dropping_groups: A utility to create a sequence of datasets $S_1, \ldots, S_{10}$ where each $S_i$ contains samples from an increasing number of categories. This is used to test how diversity metrics respond to the addition of new modes (classes).

# Example: Creating mode dropping groups
# label_to_examples is a dict mapping category names to lists of Example objects
# cats is a list of category names to include
mode_dropping_groups = data_utils.mode_dropping_groups(label_to_examples, cats, N=100)

If you have already computed a positive semi-definite similarity matrix $K$, use vendi.score_K(K) to obtain the score directly.

import numpy as np
from vendi_score import vendi

K = np.array([[1.0, 0.9, 0.0],
              [0.9, 1.0, 0.0],
              [0.0, 0.0, 1.0]])

score = vendi.score_K(K)
# Output: 2.1573

The vendi_score.text_utils module provides utilities for text diversity:

• ngram_vendi_score(sents, ns=[...]): Computes score based on n-grams. ns is a list of n-gram orders.
• embedding_vendi_score(sents, model_path="..."): Computes score using embeddings from a specified model path (e.g., BERT or SimCSE).

from vendi_score import text_utils

sents = ["Look, Jane.", "See Spot.", "See Spot run.", "Run, Spot, run.", "Jane sees Spot run."]

# N-gram based
ngram_vs = text_utils.ngram_vendi_score(sents, ns=[1, 2])

# BERT based
bert_vs = text_utils.embedding_vendi_score(sents, model_path="bert-base-uncased")

# SimCSE based
simcse_vs = text_utils.embedding_vendi_score(sents, model_path="princeton-nlp/unsup-simcse-bert-base-uncased")

When your similarity function is a dot product between normalized embeddings $X \in \mathbb{R}^{n \times d}$ and $d < n$, it is more efficient to use vendi.score_dual(X).

If the rows of $X$ are not normalized, pass normalize=True.

vendi.score_dual(X, normalize=True)

The vendi_score.image_utils module provides utilities for image diversity:

• pixel_vendi_score(imgs): Computes score based on pixel values.
• embedding_vendi_score(imgs, device="..."): Computes score using embeddings from the pool-2048 layer of the torchvision Inception v3 model.

from torchvision import datasets
from vendi_score import image_utils

mnist = datasets.MNIST("data/mnist", train=False, download=True)
digits = [[x for x, y in mnist if y == c] for c in range(10)]

# Pixel-based
pixel_vs = [image_utils.pixel_vendi_score(imgs) for imgs in digits]

# Embedding-based (Inception v3)
inception_vs = [image_utils.embedding_vendi_score(imgs, device="cuda") for imgs in digits]

Use vendi.score(samples, k, q=1.0) to compute the diversity score.

• samples: A list of elements to evaluate.
• k: A symmetric similarity function mapping a pair of elements to a score, where k(x, x) = 1.
• q: The order of the score. Large orders place greater emphasis on common elements.

import numpy as np
from vendi_score import vendi

samples = [0, 0, 10, 10, 20, 20]
k = lambda a, b: np.exp(-np.abs(a - b))

score = vendi.score(samples, k)
# Output: 2.9999

To evaluate the diversity of a set of molecules, you first need to compute a similarity kernel (matrix) using molecule_utils.get_tanimoto_K. Once you have the kernel K, you can calculate the Vendi Score using vendi.score_K(K) and the Integer Diversity using vendi.intdiv_K(K).

In the context of molecules, the fp="morgan" argument in get_tanimoto_K is commonly used to generate Morgan fingerprints for the similarity calculation.

from vendi_score import vendi, molecule_utils

# 1. Compute the Tanimoto similarity kernel using Morgan fingerprints
# Assuming 'examples' is a list of loaded molecule objects
K = molecule_utils.get_tanimoto_K([e.x for e in group.examples], fp="morgan")

# 2. Calculate diversity metrics
v_score = vendi.score_K(K)
int_div = vendi.intdiv_K(K)

print(f"Vendi Score: {v_score}")
print(f"Integer Diversity: {int_div}")

The intdiv function calculates the intrinsic diversity, which is a complementary metric to the Vendi Score.

Parameters:

• samples: An iterable of samples.
• k: A callable similarity function k(sample_i, sample_j) -> float.
• q: The order (default is 1).
• p: Optional weights for each sample.

Alternatively, use intdiv_K(K, ...) for a similarity matrix or intdiv_X(X, ...) for feature vectors.

from vendi_score.vendi import intdiv

def my_similarity(a, b):
    return 1.0 if a == b else 0.0

samples = [[1], [2], [3]]
div_val = intdiv(samples, k=my_similarity)
print(f"Intrinsic Diversity: {div_val}")

Computes the Vendi Score based on N-gram similarity. It calculates the similarity matrix $K$ by normalizing the N-gram counts and then applies vendi.score_K(K).

Parameters:

• sents: List of strings.
• ns: A list of N-gram sizes to average over (default: [1, 2, 3, 4]).
• tokenizer: Custom tokenizer.
• **kwargs: Passed to get_ngrams.

from vendi_score.text_utils import ngram_vendi_score

sents = ["Text one", "Text two", "Text three"]
score = ngram_vendi_score(sents, ns=[1, 2])