mteb 2.1.4__py3-none-any.whl → 2.5.2__py3-none-any.whl

mteb/__init__.py

@@ -9,8 +9,10 @@ from mteb.filter_tasks import filter_tasks
 from mteb.get_tasks import get_task, get_tasks
 from mteb.load_results import load_results
 from mteb.models import (
+    CacheBackendProtocol,
     CrossEncoderProtocol,
     EncoderProtocol,
+    IndexEncoderSearchProtocol,
     SearchProtocol,
     SentenceTransformerEncoderWrapper,
 )
@@ -27,8 +29,10 @@ __all__ = [
     "AbsTask",
     "Benchmark",
     "BenchmarkResults",
+    "CacheBackendProtocol",
     "CrossEncoderProtocol",
     "EncoderProtocol",
+    "IndexEncoderSearchProtocol",
     "SearchProtocol",
     "SentenceTransformerEncoderWrapper",
     "TaskMetadata",

mteb/_create_dataloaders.py

@@ -3,7 +3,7 @@ from collections.abc import Callable
 from typing import Any, cast
 
 import torch
-from datasets import Dataset
+from datasets import Dataset, Image
 from torch.utils.data import DataLoader, default_collate
 
 from mteb.abstasks.task_metadata import TaskMetadata
@@ -22,12 +22,14 @@ logger = logging.getLogger(__name__)
 def _create_dataloader_from_texts(
     text: list[str],
     batch_size: int = 32,
+    **kwargs: dict[str, Any],
 ) -> DataLoader[TextInput]:
     """Create a dataloader from a list of text.
 
     Args:
         text: A list of text to create a dataloader from.
         batch_size: Batch size for the dataloader.
+        kwargs: Not used, present catching extra arguments.
 
     Returns:
         A dataloader with the text.
@@ -244,14 +246,15 @@ def _prepare_image_dataset(
     transform: Callable[[Any], Any] | None = None,
 ) -> Dataset:
     """Prepare the image dataset by converting images to RGB and applying transformations."""
-    # If the dataset uses a different column name for images, rename it to "image".
     if (
         image_column_name
         and image_column_name in dataset.column_names
         and "image" not in dataset.column_names
     ):
         dataset = dataset.rename_column(image_column_name, "image")
-    # Map the conversion function over the dataset.
+    # don't process image if it's already in the correct format
+    if isinstance(dataset.features["image"], Image):
+        return dataset
     return dataset.map(
         _convert_images_to_rgb,
         fn_kwargs={"image_col_name": "image", "transform": transform},

mteb/_evaluators/any_sts_evaluator.py

@@ -12,6 +12,7 @@ from mteb._create_dataloaders import create_dataloader
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import compute_pairwise_similarity
+from mteb.types import PromptType
 
 from .evaluator import Evaluator
 
@@ -42,22 +43,18 @@ class AnySTSEvaluator(Evaluator):
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
+        input1_prompt_type: PromptType | None,
+        input2_prompt_type: PromptType | None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
-        self.first_column = create_dataloader(
-            dataset,
-            task_metadata,
-            input_column=sentences_column_names[0],
-        )
-        self.second_column = create_dataloader(
-            dataset,
-            task_metadata,
-            input_column=sentences_column_names[1],
-        )
+        self.dataset = dataset
+        self.input_columns = sentences_column_names
         self.task_metadata = task_metadata
         self.hf_split = hf_split
         self.hf_subset = hf_subset
+        self.input1_prompt_type = input1_prompt_type
+        self.input2_prompt_type = input2_prompt_type
 
     def __call__(
         self,
@@ -67,19 +64,31 @@ class AnySTSEvaluator(Evaluator):
     ) -> STSEvaluatorScores:
         logger.info("Running semantic similarity - Encoding samples (1/2)")
         embeddings1 = model.encode(
-            self.first_column,
+            create_dataloader(
+                self.dataset,
+                self.task_metadata,
+                input_column=self.input_columns[0],
+                **encode_kwargs,
+            ),
             task_metadata=self.task_metadata,
             hf_split=self.hf_split,
             hf_subset=self.hf_subset,
+            prompt_type=self.input1_prompt_type,
             **encode_kwargs,
         )
 
         logger.info("Running semantic similarity - Encoding samples (2/2)...")
         embeddings2 = model.encode(
-            self.second_column,
+            create_dataloader(
+                self.dataset,
+                self.task_metadata,
+                input_column=self.input_columns[1],
+                **encode_kwargs,
+            ),
             task_metadata=self.task_metadata,
             hf_split=self.hf_split,
             hf_subset=self.hf_subset,
+            prompt_type=self.input2_prompt_type,
             **encode_kwargs,
         )
 

mteb/_evaluators/classification_metrics.py

@@ -0,0 +1,54 @@
+import numpy as np
+
+
+def hamming_score(y_true: np.ndarray, y_pred: np.ndarray) -> float:
+    """Compute the Hamming score (a.k.a. label-based accuracy) for multilabel classification.
+
+    The Hamming score is the fraction of labels that are correctly predicted for each sample,
+    averaged over all samples. For samples where both y_true and y_pred have no labels,
+    the score is 1.0 (perfect agreement).
+
+    Args:
+        y_true: Binary matrix of true labels with shape (n_samples, n_labels)
+        y_pred: Binary matrix of predicted labels with shape (n_samples, n_labels)
+
+    Returns:
+        float: Hamming score between 0.0 and 1.0
+
+    Raises:
+        ValueError: If inputs are invalid or have incompatible shapes
+        TypeError: If inputs cannot be converted to numpy arrays
+    """
+    y_true = np.asarray(y_true)
+    y_pred = np.asarray(y_pred)
+
+    # Check shapes
+    if y_true.shape != y_pred.shape:
+        raise ValueError(
+            f"Shape mismatch: y_true {y_true.shape} != y_pred {y_pred.shape}"
+        )
+
+    # Check if arrays are empty
+    if y_true.size == 0:
+        raise ValueError("Input arrays cannot be empty")
+
+    # Ensure 2D arrays
+    if y_true.ndim != 2:
+        raise ValueError(f"Arrays must be 2D, got {y_true.ndim}D")
+
+    # Check for binary values
+    if not (np.all(np.isin(y_true, [0, 1])) and np.all(np.isin(y_pred, [0, 1]))):
+        raise ValueError("Arrays must contain only binary values (0 and 1)")
+
+    # Convert to boolean for bitwise operations
+    y_true_bool = y_true.astype(bool)
+    y_pred_bool = y_pred.astype(bool)
+
+    # Calculate intersection and union for each sample
+    intersection = (y_true_bool & y_pred_bool).sum(axis=1)
+    union = (y_true_bool | y_pred_bool).sum(axis=1)
+
+    # Handle division by zero: when union is 0, both are all zeros, so score is 1.0
+    scores = np.where(union == 0, 1.0, intersection / union)
+
+    return float(scores.mean())

mteb/_evaluators/clustering_evaluator.py

@@ -44,7 +44,7 @@ class ClusteringEvaluator(Evaluator):
             self.dataset,
             self.task_metadata,
             input_column=self.input_column_name,
-            batch_size=encode_kwargs["batch_size"],
+            **encode_kwargs,
         )
 
         logger.info("Running clustering - Encoding samples...")

mteb/_evaluators/image/imagetext_pairclassification_evaluator.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.nn.functional as F
 from datasets import Dataset
-from PIL.Image import Image
 from torch.utils.data import DataLoader
 
 from mteb._create_dataloaders import (
@@ -15,6 +16,10 @@ from mteb._requires_package import requires_image_dependencies
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models.models_protocols import EncoderProtocol
 
+if TYPE_CHECKING:
+    from PIL.Image import Image
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -103,7 +108,7 @@ class ImageTextPairClassificationEvaluator(Evaluator):
         text_embeddings = model.encode(
             DataLoader(
                 Dataset.from_dict({"text": texts}),
-                batch_size=encode_kwargs["batch_size"],
+                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,
@@ -122,8 +127,8 @@ class ImageTextPairClassificationEvaluator(Evaluator):
         image_embeddings = model.encode(
             DataLoader(
                 CustomImageDataset(images),
-                batch_size=encode_kwargs["batch_size"],
                 collate_fn=lambda x: {"image": [item["image"] for item in x]},
+                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,

mteb/_evaluators/pair_classification_evaluator.py

@@ -14,6 +14,7 @@ from mteb._evaluators.evaluator import Evaluator
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.similarity_functions import compute_pairwise_similarity
+from mteb.types import PromptType
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +61,8 @@ class PairClassificationEvaluator(Evaluator):
         task_metadata: TaskMetadata,
         hf_split: str,
         hf_subset: str,
+        input1_prompt_type: PromptType | None,
+        input2_prompt_type: PromptType | None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -69,6 +72,8 @@ class PairClassificationEvaluator(Evaluator):
         self.task_metadata = task_metadata
         self.hf_split = hf_split
         self.hf_subset = hf_subset
+        self.input1_prompt_type = input1_prompt_type
+        self.input2_prompt_type = input2_prompt_type
 
         if len(self.dataset[self.input1_column_name]) != len(
             self.dataset[self.input2_column_name]
@@ -82,47 +87,34 @@ class PairClassificationEvaluator(Evaluator):
         model: EncoderProtocol,
         encode_kwargs: dict[str, Any],
     ) -> PairClassificationDistances:
-        logger.info("Running pair classification - Encoding inputs...")
-        if self.task_metadata.modalities == ["text"]:
-            # datasets v4 will pass column objects, so we need to extract the text
-            all_sentences = (
-                self.dataset[self.input1_column_name][:]
-                + self.dataset[self.input2_column_name][:]
-            )
-            len_sentences1 = len(self.dataset[self.input1_column_name])
-            embeddings = self._encode_unique_texts(
-                all_sentences,
-                model,
-                task_metadata=self.task_metadata,
-                hf_split=self.hf_split,
-                hf_subset=self.hf_subset,
-                **encode_kwargs,
-            )
-            embeddings1 = embeddings[:len_sentences1]
-            embeddings2 = embeddings[len_sentences1:]
-        else:
-            embeddings1 = model.encode(
-                create_dataloader(
-                    self.dataset,
-                    task_metadata=self.task_metadata,
-                    input_column=self.input1_column_name,
-                ),
+        logger.info("Running pair classification - Encoding samples (1/2)")
+        embeddings1 = model.encode(
+            create_dataloader(
+                self.dataset,
                 task_metadata=self.task_metadata,
-                hf_split=self.hf_split,
-                hf_subset=self.hf_subset,
+                input_column=self.input1_column_name,
                 **encode_kwargs,
-            )
-            embeddings2 = model.encode(
-                create_dataloader(
-                    self.dataset,
-                    task_metadata=self.task_metadata,
-                    input_column=self.input2_column_name,
-                ),
+            ),
+            task_metadata=self.task_metadata,
+            hf_split=self.hf_split,
+            hf_subset=self.hf_subset,
+            prompt_type=self.input1_prompt_type,
+            **encode_kwargs,
+        )
+        logger.info("Running pair classification - Encoding samples (2/2)")
+        embeddings2 = model.encode(
+            create_dataloader(
+                self.dataset,
                 task_metadata=self.task_metadata,
-                hf_split=self.hf_split,
-                hf_subset=self.hf_subset,
+                input_column=self.input2_column_name,
                 **encode_kwargs,
-            )
+            ),
+            task_metadata=self.task_metadata,
+            hf_split=self.hf_split,
+            hf_subset=self.hf_subset,
+            prompt_type=self.input2_prompt_type,
+            **encode_kwargs,
+        )
 
         logger.info("Running pair classification - Evaluating pair similarity...")
         cosine_scores = 1 - paired_cosine_distances(embeddings1, embeddings2)
@@ -168,7 +160,7 @@ class PairClassificationEvaluator(Evaluator):
         )
         all_unique_texts_embs = np.asarray(
             model.encode(
-                _create_dataloader_from_texts(all_unique_texts),
+                _create_dataloader_from_texts(all_unique_texts, **encode_kwargs),
                 task_metadata=task_metadata,
                 hf_split=hf_split,
                 hf_subset=hf_subset,

mteb/_evaluators/sklearn_evaluator.py

@@ -6,7 +6,7 @@ from datasets import Dataset
 from torch.utils.data import DataLoader
 from typing_extensions import Self
 
-from mteb._create_dataloaders import _create_image_dataloader
+from mteb._create_dataloaders import create_dataloader
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.models import EncoderProtocol
 from mteb.types import BatchedInput
@@ -50,33 +50,20 @@ class SklearnEvaluator(Evaluator):
         self.evaluator_model = evaluator_model
 
     def create_dataloaders(
-        self, batch_size: int
+        self, encode_kwargs: dict[str, Any]
     ) -> tuple[DataLoader[BatchedInput], DataLoader[BatchedInput]]:
-        if self.task_metadata.modalities == ["image"]:
-            dataloader_train = _create_image_dataloader(
-                self.train_dataset,
-                image_column_name=self.values_column_name,
-                batch_size=batch_size,
-            )
-            dataloader_test = _create_image_dataloader(
-                self.eval_dataset,
-                image_column_name=self.values_column_name,
-                batch_size=batch_size,
-            )
-        elif self.task_metadata.modalities == ["text"]:
-            if self.values_column_name != "text":
-                self.train_dataset = self.train_dataset.rename_column(
-                    self.values_column_name, "text"
-                )
-                self.eval_dataset = self.eval_dataset.rename_column(
-                    self.values_column_name, "text"
-                )
-            dataloader_train = DataLoader(self.train_dataset)
-            dataloader_test = DataLoader(self.eval_dataset)
-        else:
-            raise ValueError(
-                "ClassificationEvaluator only supports image and text modalities."
-            )
+        dataloader_train = create_dataloader(
+            self.train_dataset,
+            self.task_metadata,
+            input_column=self.values_column_name,
+            **encode_kwargs,
+        )
+        dataloader_test = create_dataloader(
+            self.eval_dataset,
+            self.task_metadata,
+            input_column=self.values_column_name,
+            **encode_kwargs,
+        )
         return dataloader_train, dataloader_test
 
     def __call__(  # type: ignore[override]
@@ -98,7 +85,7 @@ class SklearnEvaluator(Evaluator):
 
         """
         dataloader_train, dataloader_test = self.create_dataloaders(
-            batch_size=encode_kwargs["batch_size"]
+            encode_kwargs=encode_kwargs,
         )
 
         logger.info("Running - Encoding samples...")

mteb/_evaluators/text/bitext_mining_evaluator.py

@@ -46,7 +46,10 @@ class BitextMiningEvaluator(Evaluator):
 
         embeddings = {}
         for sub in tqdm(subsets):
-            dataloader = _create_dataloader_from_texts(self.sentences[sub])
+            dataloader = _create_dataloader_from_texts(
+                self.sentences[sub],
+                **encode_kwargs,
+            )
             embeddings[sub] = model.encode(
                 dataloader,
                 task_metadata=self.task_metadata,

mteb/_evaluators/text/summarization_evaluator.py

@@ -109,7 +109,8 @@ class SummarizationEvaluator(Evaluator):
                     summary
                     for human_summaries in self.human_summaries
                     for summary in human_summaries
-                ]
+                ],
+                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,
@@ -124,7 +125,8 @@ class SummarizationEvaluator(Evaluator):
                     summary
                     for machine_summaries in self.machine_summaries
                     for summary in machine_summaries
-                ]
+                ],
+                **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -42,14 +42,14 @@ class ZeroShotClassificationEvaluator(Evaluator):
     ) -> Array:
         dataloader = create_dataloader(
             self.dataset,
-            batch_size=encode_kwargs["batch_size"],
             input_column=self.input_column_name,
             task_metadata=self.task_metadata,
+            **encode_kwargs,
         )
 
         logger.info("Running zero-shot classification - Encoding labels...")
         text_label_embeddings = model.encode(
-            _create_dataloader_from_texts(self.candidate_labels),
+            _create_dataloader_from_texts(self.candidate_labels, **encode_kwargs),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,
             hf_split=self.hf_split,

mteb/abstasks/_data_filter/filters.py

@@ -0,0 +1,125 @@
+"""Simplified version of https://gist.github.com/AlexeyVatolin/ea3adc21aa7a767603ff393b22085adc from https://github.com/embeddings-benchmark/mteb/pull/2900"""
+
+import logging
+
+import datasets
+import pandas as pd
+from datasets import Dataset, DatasetDict
+
+from mteb import TaskMetadata
+
+logger = logging.getLogger(__name__)
+
+
+def deduplicate(dataset: Dataset, input_column: str) -> Dataset:
+    """Remove duplicate texts, keeping the first occurrence."""
+    unique_texts = set()
+    indices_to_keep = []
+    for i, text in enumerate(dataset[input_column]):
+        text = text.strip()
+        if text not in unique_texts:
+            unique_texts.add(text)
+            indices_to_keep.append(i)
+
+    logger.info(
+        f"[deduplicate] removed={len(dataset) - len(indices_to_keep)}/{len(dataset)}"
+    )
+    return dataset.select(indices_to_keep)
+
+
+def filter_empty(dataset: Dataset, input_column: str) -> Dataset:
+    """Filter out empty or whitespace-only examples."""
+    before = len(dataset)
+    ds = dataset.filter(lambda x: len(x[input_column].strip()) > 0)
+    logger.info(f"[filter_empty] removed={before - len(ds)}/{before}")
+    return ds
+
+
+def filter_train_leakage(
+    train_dataset: Dataset, test_dataset: Dataset, input_column: str
+) -> Dataset:
+    """Remove test examples that appear in training."""
+    train_texts = set(train_dataset[input_column])
+    before = len(test_dataset)
+    indices = [
+        i
+        for i, text in enumerate(test_dataset[input_column])
+        if text not in train_texts
+    ]
+    logger.info(f"[filter_train_leakage] removed={before - len(indices)}/{before}")
+    return test_dataset.select(indices)
+
+
+def filter_unclear_label(
+    dataset_dict: DatasetDict, input_column: str, label_column: str
+) -> DatasetDict:
+    """Remove examples where the same text appears with multiple different labels."""
+    normalized: dict[str, set[str | tuple[str, ...]]] = {}
+    logger.debug("[filter_controversial] scanning dataset for label conflicts...")
+
+    for split, ds in dataset_dict.items():
+        for text, label in zip(ds[input_column], ds[label_column]):
+            key = text.strip().lower()
+            normalized.setdefault(key, set()).add(
+                label if isinstance(label, (str, int, float)) else tuple(label)
+            )
+
+    bad_texts = {t for t, labels in normalized.items() if len(labels) > 1}
+    logger.info(f"[filter_controversial] Removing {len(bad_texts)} conflicting texts")
+
+    new_dict = {}
+    for split, ds in dataset_dict.items():
+        before = len(ds)
+        filtered = ds.filter(lambda x: x[input_column].strip().lower() not in bad_texts)
+        logger.debug(
+            f"[filter_controversial:{split}] removed={before - len(filtered)}/{before}"
+        )
+        new_dict[split] = filtered
+
+    return DatasetDict(new_dict)
+
+
+def filter_short(dataset: Dataset, input_column: str, min_words: int = 3) -> Dataset:
+    """Filter out texts with fewer than `min_words`."""
+    before = len(dataset)
+    ds = dataset.filter(lambda x: len(x[input_column].strip().split()) >= min_words)
+    logger.debug(f"[filter_short] removed={before - len(ds)}/{before}")
+    return ds
+
+
+def split_train_test(
+    ds: DatasetDict,
+    metadata: TaskMetadata,
+    train_split: str,
+    label_column: str,
+) -> DatasetDict:
+    if train_split in ds and metadata.eval_splits == train_split:
+        before = len(ds[train_split])
+        logger.info(
+            f"[split_train_test] eval_splits == train_split; performing split on {before} examples"
+        )
+        ds[train_split] = ds[train_split].cast_column(
+            label_column,
+            datasets.ClassLabel(names=list(set(ds[train_split][label_column]))),
+        )
+        label_counts = pd.Series(ds[train_split][label_column]).value_counts()
+        one_sample_labels = set(label_counts[label_counts == 1].index.tolist())
+
+        if one_sample_labels:
+            logger.info(
+                f"[split_train_test] Removing {len(one_sample_labels)} labels with only one instance"
+            )
+            ds[train_split] = ds[train_split].filter(
+                lambda x: x[label_column] not in one_sample_labels
+            )
+
+        splits = ds[train_split].train_test_split(
+            test_size=min(2048, before // 2), seed=42, stratify_by_column=label_column
+        )
+        ds = DatasetDict({train_split: splits[train_split], "test": splits["test"]})
+        metadata.eval_splits = ["test"]
+        logger.info(
+            f"[split_train_test] Train size={len(ds[train_split])}, Test size={len(ds['test'])}"
+        )
+
+    return ds