Evaluation

sentence_transformers.multi_vector_encoder.evaluation defines evaluators tailored to multi-vector (late-interaction) models. All evaluators score with MaxSim and accept the multi-vector model’s ragged per-token embeddings.

MultiVectorInformationRetrievalEvaluator

class sentence_transformers.multi_vector_encoder.evaluation.MultiVectorInformationRetrievalEvaluator(queries: dict[str, SingleInput], corpus: dict[str, SingleInput], relevant_docs: dict[str, set[str]], *, corpus_chunk_size: int = 5000, chunk_elements: int | None = None, score_functions: dict[str, Callable[[Tensor, Tensor], Tensor]] | None = None, **kwargs)[source]

Evaluates a MultiVectorEncoder model on an information-retrieval (IR) task. For each query, the model retrieves the top-k closest documents from the corpus using ColBERT-style MaxSim scoring, then reports the standard IR metrics (MRR@k, NDCG@k, Recall@k, Precision@k, Accuracy@k, MAP@k) against the supplied relevance judgements.

Parameters:

queries (Dict[str, str]) – Mapping of query IDs to query text.
corpus (Dict[str, str]) – Mapping of document IDs to document text.
relevant_docs (Dict[str, Set[str]]) – Mapping of query ID to the set of relevant document IDs.
corpus_chunk_size (int) – How many documents to encode and score per round-trip. Larger values mean more encoded doc embeddings live in memory at once but fewer encode-pass round-trips. Defaults to 5000.
chunk_elements (int, optional) – Element budget for the 4D (batch_q, chunk, q_tokens, d_tokens) MaxSim scoring intermediate, forwarded to maxsim(), which packs document chunks under it, adapting to the query count and document lengths. Defaults to None (maxsim’s 100M-element budget, at most ~400 MB, half that in bf16 / fp16). Lower it to cut evaluation memory.
score_functions (Dict[str, Callable], optional) – Override the default scoring, which resolves from the model’s similarity_fn_name at call time (with chunk_elements applied if one was given). The chosen callable receives (queries, documents) token tensors and must return a (num_queries, num_documents) score matrix. XTR scoring is not supported here because it does a global top-k across the whole candidate set, which is incompatible with this evaluator’s per-chunk corpus scoring.
mrr_at_k (List[int]) – k-values for MRR. Defaults to [10].
ndcg_at_k (List[int]) – k-values for NDCG. Defaults to [10].
accuracy_at_k (List[int]) – k-values for accuracy. Defaults to [1, 3, 5, 10].
precision_recall_at_k (List[int]) – k-values for precision and recall. Defaults to [1, 3, 5, 10].
map_at_k (List[int]) – k-values for MAP. Defaults to [100].
show_progress_bar (bool) – Show a progress bar during evaluation. Defaults to False.
batch_size (int) – Per-input batch size used while encoding. Defaults to 32.
name (str) – Evaluation name (used as the dataset stem in CSV / prediction filenames). Defaults to "".
write_csv (bool) – Append per-call metric values to Information-Retrieval_evaluation_<name>_results.csv. Defaults to True.
truncate_dim (int, optional) – Not supported: multi-vector token embeddings have no Matryoshka-style truncation, so any non-None value raises a ValueError. Defaults to None.
main_score_function (str or SimilarityFunction, optional) – Which score-function key to treat as the primary metric for the model card / trainer. Defaults to None (use the first / only key in score_functions).
query_prompt (str, optional) – Prompt prepended to every query during encoding. Defaults to None.
query_prompt_name (str, optional) – Name of a prompt registered on the model to prepend to queries. Mutually exclusive with query_prompt. Defaults to None.
corpus_prompt (str, optional) – Prompt prepended to every corpus document. Defaults to None.
corpus_prompt_name (str, optional) – Name of a prompt registered on the model to prepend to corpus documents. Mutually exclusive with corpus_prompt. Defaults to None.
write_predictions (bool) – Write per-query top-k predictions to a JSONL file, suitable as input to ReciprocalRankFusionEvaluator. Defaults to False.

Example

from datasets import load_dataset

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorInformationRetrievalEvaluator

model = MultiVectorEncoder("lightonai/GTE-ModernColBERT-v1")

# Load NanoMSMARCO subsets and convert to the evaluator's dict format.
corpus_ds = load_dataset("sentence-transformers/NanoBEIR-en", "corpus", split="NanoMSMARCO")
queries_ds = load_dataset("sentence-transformers/NanoBEIR-en", "queries", split="NanoMSMARCO")
qrels_ds = load_dataset("sentence-transformers/NanoBEIR-en", "qrels", split="NanoMSMARCO")

corpus = {row["_id"]: row["text"] for row in corpus_ds}
queries = {row["_id"]: row["text"] for row in queries_ds}
relevant_docs = {}
for row in qrels_ds:
    relevant_docs.setdefault(row["query-id"], set()).add(row["corpus-id"])

evaluator = MultiVectorInformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="NanoMSMARCO",
)
results = evaluator(model)
print(results[evaluator.primary_metric])

MultiVectorNanoBEIREvaluator

class sentence_transformers.multi_vector_encoder.evaluation.MultiVectorNanoBEIREvaluator(*args, corpus_chunk_size: int = 5000, chunk_elements: int | None = None, **kwargs)[source]

Evaluates a MultiVectorEncoder model on the NanoBEIR collection.

NanoBEIR is a downsized version of BEIR (around 50 queries and 5,000 documents per subset) used for quick retrieval-quality benchmarking before running a full-scale BEIR evaluation. This evaluator runs a MultiVectorInformationRetrievalEvaluator on each requested Nano-* subset, reports the same IR metrics per dataset (MRR@k, NDCG@k, Recall@k, Precision@k, Accuracy@k, MAP@k), and aggregates them across datasets at the end.

Parameters:

dataset_names (List[str], optional) – Short names of NanoBEIR subsets to evaluate ("climatefever", "dbpedia", "fever", "fiqa2018", "hotpotqa", "msmarco", "nfcorpus", "nq", "quoraretrieval", "scidocs", "arguana", "scifact", "touche2020"). Defaults to every subset.
dataset_id (str) – The HuggingFace dataset ID hosting the corpus / queries / qrels subsets. Defaults to "sentence-transformers/NanoBEIR-en". Swap in a translated variant from the NanoBEIR collection for non-English evaluation.
corpus_chunk_size (int) – How many documents to encode + score per round-trip. Larger values mean more encoded doc embeddings live in memory at once but fewer encode-pass round-trips. Defaults to 5000.
chunk_elements (int, optional) – Element budget for the 4D (batch_q, chunk, q_tokens, d_tokens) MaxSim scoring intermediate, forwarded to maxsim(), which packs document chunks under it, adapting to the query count and document lengths. Defaults to None (maxsim’s 100M-element budget, at most ~400 MB, half that in bf16 / fp16). Lower it to cut evaluation memory.
mrr_at_k (List[int]) – k-values for MRR. Defaults to [10].
ndcg_at_k (List[int]) – k-values for NDCG. Defaults to [10].
accuracy_at_k (List[int]) – k-values for accuracy. Defaults to [1, 3, 5, 10].
precision_recall_at_k (List[int]) – k-values for precision and recall. Defaults to [1, 3, 5, 10].
map_at_k (List[int]) – k-values for MAP. Defaults to [100].
show_progress_bar (bool) – Show a progress bar during evaluation. Defaults to False.
batch_size (int) – Per-input batch size used while encoding. Defaults to 32.
write_csv (bool) – Append per-call metric values to a CSV file (one row per evaluation call). Defaults to True.
truncate_dim (int, optional) – Not supported: multi-vector token embeddings have no Matryoshka-style truncation, so any non-None value raises a ValueError. Defaults to None.
score_functions (Dict[str, Callable], optional) – Override the default per-subset scoring. See MultiVectorInformationRetrievalEvaluator. Defaults to None.
main_score_function (str or SimilarityFunction, optional) – Score-function key to treat as the primary metric for the model card / trainer. Defaults to None.
aggregate_fn (Callable[[List[float]], float]) – How to aggregate the per-subset scores into the top-level summary. Defaults to np.mean.
aggregate_key (str) – Suffix used in the aggregated metric’s key. Defaults to "mean".
query_prompts (str or Dict[str, str], optional) – Prompt prepended to queries. A string applies to every dataset. A dict keyed by dataset_names entries lets you set per-dataset prompts. Defaults to None.
corpus_prompts (str or Dict[str, str], optional) – Prompt prepended to corpus documents. Same string / dict semantics as query_prompts. Defaults to None.
write_predictions (bool) – Write per-query top-k predictions to a JSONL file, suitable as input to ReciprocalRankFusionEvaluator. Defaults to False.

Example

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator

model = MultiVectorEncoder("lightonai/GTE-ModernColBERT-v1")

evaluator = MultiVectorNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus"])
results = evaluator(model)
print(results[evaluator.primary_metric])

MultiVectorTripletEvaluator

class sentence_transformers.multi_vector_encoder.evaluation.MultiVectorTripletEvaluator(*args, **kwargs)[source]

Triplet evaluator for MultiVectorEncoder models.

Given (anchor, positive, negative) triplets, checks how often MaxSim(anchor, positive) > MaxSim(anchor, negative) + margin. The anchors are encoded via encode_query() (with the query prefix and length). Positives and negatives are encoded via encode_document(). truncate_dim is not supported (multi-vector token embeddings have no Matryoshka-style truncation) and raises a ValueError.

A margin dictionary must be keyed by a supported similarity, maxsim or meanmaxsim (a float applies to both). The two need different margins: MaxSim sums a per-token maximum over the query tokens, so scores grow with the query length (roughly one point per query token for normalized embeddings), while MeanMaxSim divides that by the token count and sits in cosine’s [-1, 1] range.

Example

from datasets import load_dataset

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorTripletEvaluator

model = MultiVectorEncoder("lightonai/GTE-ModernColBERT-v1")

dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")

evaluator = MultiVectorTripletEvaluator(
    anchors=dataset[:1000]["anchor"],
    positives=dataset[:1000]["positive"],
    negatives=dataset[:1000]["negative"],
    name="all-nli-dev",
)
results = evaluator(model)
print(results[evaluator.primary_metric])

MultiVectorRerankingEvaluator

class sentence_transformers.multi_vector_encoder.evaluation.MultiVectorRerankingEvaluator(samples: list[dict], *, at_k: int = 10, name: str = '', write_csv: bool = True, similarity_fct: Callable[[Tensor | list[Tensor], list[Tensor]], Tensor] | None = None, batch_size: int = 16, show_progress_bar: bool = False, use_batched_encoding: bool = True)[source]

Reranking evaluator for MultiVectorEncoder models.

Scores each query’s candidate documents with the model’s multi-vector similarity (MaxSim by default) and reports MAP, MRR@k, and NDCG@k, treating the positive documents as the relevance ground truth. Useful for evaluating a multi-vector model as a second-stage reranker: a first-stage retriever returns candidates per query (positives mixed with distractors) and the multi-vector model rescores them.

See RerankingEvaluator for the full argument list and the samples format. This subclass differs only by resolving the scoring from the model at call time and encoding queries / documents asymmetrically (via encode_query / encode_document).

Example

from datasets import load_dataset

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorRerankingEvaluator

model = MultiVectorEncoder("lightonai/GTE-ModernColBERT-v1")

# Each sample pairs a query with its relevant documents and the first-stage distractors.
eval_dataset = load_dataset("microsoft/ms_marco", "v1.1", split="validation").select(range(100))
samples = [
    {
        "query": sample["query"],
        "positive": [
            text
            for selected, text in zip(sample["passages"]["is_selected"], sample["passages"]["passage_text"])
            if selected
        ],
        "negative": [
            text
            for selected, text in zip(sample["passages"]["is_selected"], sample["passages"]["passage_text"])
            if not selected
        ],
    }
    for sample in eval_dataset
]

evaluator = MultiVectorRerankingEvaluator(samples=samples, name="ms-marco-dev")
results = evaluator(model)
print(results[evaluator.primary_metric])

MultiVectorDistillationEvaluator

class sentence_transformers.multi_vector_encoder.evaluation.MultiVectorDistillationEvaluator(queries: Sequence[SingleInput], documents: Sequence[SingleInput] | Sequence[Sequence[SingleInput]], scores: list[float] | list[list[float]] | torch.Tensor, *, temperature: float = 1.0, student_temperature: float | None = None, teacher_temperature: float | None = None, similarity_fct: Callable | None = None, name: str = '', batch_size: int = 16, show_progress_bar: bool = False, write_csv: bool = True)[source]

Distillation evaluator for MultiVectorEncoder models.

Two data shapes are supported:

Per-query candidate sets (the KD training format, matching MultiVectorDistillKLDivLoss and PyLate): documents is a list of N-way candidate lists per query and scores the matching 2-D teacher scores. Both metrics are computed per query, so they track the training loss: the KL divergence over each query’s own candidate set (with the same temperature handling as the loss, averaged over queries, so with the loss’s temperatures, and its scoring mirrored via similarity_fct when non-default, the reported KL matches the training loss) and the Spearman as the mean of the per-query rank correlations.
Flat pairs: one document per query with 1-D scores. A per-query distribution or ranking is undefined here, so the KL softmaxes over the whole dataset as a single distribution and reports the full divergence (a sum over pairs, not divided by their number, and not comparable to the per-query KL or to PyLate). The temperatures divide the scores the same way as on the per-query path. The Spearman is one global correlation over all pairs.

Reported metrics:

KL divergence between teacher and student score distributions (lower is better).
Spearman rank correlation between teacher and student scores (higher is better, the primary metric). With candidate sets this is the mean of the per-query correlations, matching the per-query KL above. It also avoids a scale trap: MaxSim sums over query tokens, so absolute scores are not comparable across queries (MeanMaxSim divides that out, but the per-query correlation is the one that tracks the loss either way) and a single correlation over all scores would mostly measure per-query offsets rather than ranking. Queries where the teacher or student scores are constant have no defined rank correlation and are skipped from the mean (0.0 is reported if every query is skipped).

The Spearman score is generally a more interpretable mid-training signal than raw KL.

Parameters:

queries – List of query texts.
documents – One document per query, or a list of N candidate documents per query (N constant across queries). Must have the same length as queries.
scores – Teacher scores: 1-D (one per pair) for flat documents, 2-D (num_queries, N) for candidate sets.
temperature – Temperature dividing both the teacher and student scores before the softmax unless overridden per side, with the KL scaled by the student temperature squared, exactly as in MultiVectorDistillKLDivLoss. Set it to the training loss’s temperature so the per-query KL matches the training loss. Must be > 0. Defaults to 1.0.
student_temperature – Student-side override of temperature, mirroring the loss’s parameter of the same name. Must be > 0. Defaults to None.
teacher_temperature – Teacher-side override of temperature, mirroring the loss’s parameter of the same name. Must be > 0. Defaults to None.
similarity_fct – Pairwise scoring callable replacing model.similarity_pairwise, to mirror a non-default training similarity_fct: e.g. pass mean_colbert_scores_pairwise when the loss uses MeanMaxSim scoring, or the per-query KL cannot match the training loss. Setting model.similarity_fn_name = "meanmaxsim" covers the same case without this argument. Defaults to None (the model’s own pairwise similarity).
name – Optional run name appended to CSV filenames.
batch_size – Batch size for encoding.
show_progress_bar – Whether to show a progress bar.
write_csv – Whether to write per-call results to a CSV file under output_path.

Example

from sentence_transformers import MultiVectorEncoder
from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorDistillationEvaluator

model = MultiVectorEncoder("lightonai/GTE-ModernColBERT-v1")

queries = ["What is the capital of France?", "Who painted the Mona Lisa?"]
# One candidate list per query, with the matching teacher scores.
documents = [
    ["Paris is the capital of France.", "Berlin is the capital of Germany."],
    ["Leonardo da Vinci painted the Mona Lisa.", "Van Gogh painted The Starry Night."],
]
scores = [[9.5, 2.1], [8.8, 1.4]]

evaluator = MultiVectorDistillationEvaluator(
    queries=queries,
    documents=documents,
    scores=scores,
    # Match the training loss's temperature so the reported KL tracks the training loss.
    temperature=0.25,
    name="msmarco-dev",
)
results = evaluator(model)
print(results[evaluator.primary_metric])