mteb 2.0.5__py3-none-any.whl → 2.1.19__py3-none-any.whl

mteb/__init__.py

@@ -1,5 +1,6 @@
 from importlib.metadata import version
 
+from mteb import types
 from mteb.abstasks import AbsTask
 from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.deprecated_evaluator import MTEB
@@ -7,7 +8,12 @@ from mteb.evaluate import evaluate
 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 EncoderProtocol, SentenceTransformerEncoderWrapper
+from mteb.models import (
+    CrossEncoderProtocol,
+    EncoderProtocol,
+    SearchProtocol,
+    SentenceTransformerEncoderWrapper,
+)
 from mteb.models.get_model_meta import get_model, get_model_meta, get_model_metas
 from mteb.results import BenchmarkResults, TaskResult
 
@@ -21,7 +27,9 @@ __all__ = [
     "AbsTask",
     "Benchmark",
     "BenchmarkResults",
+    "CrossEncoderProtocol",
     "EncoderProtocol",
+    "SearchProtocol",
     "SentenceTransformerEncoderWrapper",
     "TaskMetadata",
     "TaskResult",
@@ -35,4 +43,5 @@ __all__ = [
     "get_task",
     "get_tasks",
     "load_results",
+    "types",
 ]

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},
@@ -277,6 +280,8 @@ def _custom_collate_fn(batch: list[dict[str, Any]]) -> dict[str, Any]:
             # Leave the images as a list to avoid stacking errors.
             collated[key] = [item[key] for item in batch]
         else:
+            if any(item[key] is None for item in batch):
+                raise ValueError(f"Found None in batch for key '{key}'")
             collated[key] = default_collate([item[key] for item in batch])
     return collated
 

mteb/_evaluators/any_sts_evaluator.py

@@ -45,16 +45,8 @@ class AnySTSEvaluator(Evaluator):
         **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
@@ -67,7 +59,12 @@ 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,
@@ -76,7 +73,12 @@ class AnySTSEvaluator(Evaluator):
 
         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,

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

@@ -103,7 +103,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 +122,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

@@ -106,6 +106,7 @@ class PairClassificationEvaluator(Evaluator):
                     self.dataset,
                     task_metadata=self.task_metadata,
                     input_column=self.input1_column_name,
+                    **encode_kwargs,
                 ),
                 task_metadata=self.task_metadata,
                 hf_split=self.hf_split,
@@ -117,6 +118,7 @@ class PairClassificationEvaluator(Evaluator):
                     self.dataset,
                     task_metadata=self.task_metadata,
                     input_column=self.input2_column_name,
+                    **encode_kwargs,
                 ),
                 task_metadata=self.task_metadata,
                 hf_split=self.hf_split,
@@ -168,7 +170,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/retrieval_metrics.py

@@ -5,7 +5,6 @@ from typing import Any
 import numpy as np
 import pandas as pd
 import pytrec_eval
-import torch
 from packaging.version import Version
 from sklearn.metrics import auc
 
@@ -14,14 +13,6 @@ from mteb.types import RelevantDocumentsType, RetrievalEvaluationResult
 logger = logging.getLogger(__name__)
 
 
-try:
-    # speeds up computation if available
-    torch.set_float32_matmul_precision("high")
-    logger.info("Setting torch float32 matmul precision to high for a speedup")
-except Exception:
-    pass
-
-
 def mrr(
     qrels: RelevantDocumentsType,
     results: dict[str, dict[str, float]],

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/_stratification.py

@@ -134,7 +134,7 @@ def _get_most_desired_combination(samples_with_combination: dict):
 class IterativeStratification(_BaseKFold):
     """Iteratively stratify a multi-label data set into folds
 
-    Construct an interative stratifier that splits the data set into folds trying to maintain balanced representation
+    Construct an iterative stratifier that splits the data set into folds trying to maintain balanced representation
     with respect to order-th label combinations.
     """
 

mteb/abstasks/abstask.py

@@ -459,7 +459,7 @@ class AbsTask(ABC):
         """Filter the languages of the task.
 
         Args:
-            languages: list of languages to filter the task by can be either a 3-letter langauge code (e.g. "eng") or also include the script
+            languages: list of languages to filter the task by can be either a 3-letter language code (e.g. "eng") or also include the script
                 (e.g. "eng-Latn")
             script: A list of scripts to filter the task by. Will be ignored if language code specified the script. If None, all scripts are included.
                 If the language code does not specify the script the intersection of the language and script will be used.
@@ -491,6 +491,11 @@ class AbsTask(ABC):
                 if lang_scripts.contains_languages(langs):
                     subsets_to_keep.append(hf_subset)
 
+        if len(subsets_to_keep) == 0:
+            raise ValueError(
+                f"No subsets were found for {self.metadata.name} with filters: language code {languages}, script {script}, hf subsets {hf_subsets}."
+            )
+
         self.hf_subsets = subsets_to_keep
         return self
 

mteb/abstasks/clustering.py

@@ -200,7 +200,7 @@ class AbsTaskClustering(AbsTask):
                 downsampled_dataset,
                 self.metadata,
                 input_column=self.input_column_name,
-                batch_size=encode_kwargs["batch_size"],
+                **encode_kwargs,
             ),
             task_metadata=self.metadata,
             hf_subset=hf_subset,

mteb/abstasks/dataset_card_template.md

@@ -85,7 +85,7 @@ desc_stats = task.metadata.descriptive_stats
 ```
 
 ```json
-{{ descritptive_stats | default("{}", true) }}
+{{ descriptive_stats | default("{}", true) }}
 ```
 
 </details>

mteb/abstasks/multilabel_classification.py

@@ -112,7 +112,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
             unique_train_dataset,
             self.metadata,
             input_column=self.input_column_name,
-            batch_size=encode_kwargs["batch_size"],
+            **encode_kwargs,
         )
 
         logger.info("Running multilabel classification - Encoding training set...")
@@ -141,7 +141,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
             test_dataset.select_columns(self.input_column_name),
             self.metadata,
             input_column=self.input_column_name,
-            batch_size=encode_kwargs["batch_size"],
+            **encode_kwargs,
         )
 
         logger.info("Running multilabel classification - Encoding test set...")

mteb/abstasks/retrieval.py

@@ -653,6 +653,8 @@ class AbsTaskRetrieval(AbsTask):
             FileNotFoundError: If the specified path does not exist.
             ValueError: If the loaded top ranked results are not in the expected format.
         """
+        self._top_k = top_k
+
         top_ranked_path = Path(top_ranked_path)
         if top_ranked_path.is_dir():
             top_ranked_path = self._predictions_path(top_ranked_path)
@@ -682,7 +684,6 @@ class AbsTaskRetrieval(AbsTask):
                     top_k_sorted[query_id] = sorted_keys[: self._top_k]
 
                 self.dataset[subset][split]["top_ranked"] = top_k_sorted
-        self._top_k = top_k
         return self
 
 

mteb/abstasks/retrieval_dataset_loaders.py

@@ -176,7 +176,7 @@ class RetrievalDatasetLoader:
                 {
                     "query-id": Value("string"),
                     "corpus-id": Value("string"),
-                    "score": Value("uint16"),
+                    "score": Value("int32"),
                 }
             )
         )

mteb/abstasks/task_metadata.py

@@ -107,6 +107,7 @@ The domains follow the categories used in the [Universal Dependencies project](h
 SampleCreationMethod = Literal[
     "found",
     "created",
+    "created and machine-translated",
     "human-translated and localized",
     "human-translated",
     "machine-translated",
@@ -532,7 +533,7 @@ class TaskMetadata(BaseModel):
                 citation=self.bibtex_citation,
                 dataset_description=self.description,
                 dataset_reference=self.reference,
-                descritptive_stats=descriptive_stats,
+                descriptive_stats=descriptive_stats,
                 dataset_task_name=self.name,
                 category=self.category,
                 domains=", ".join(self.domains) if self.domains else None,

mteb/benchmarks/_create_table.py

@@ -358,9 +358,7 @@ def _create_summary_table_mean_public_private(
         "mean(public)": "Mean (Public)",
         "mean(private)": "Mean (Private)",
     }
-    # For RTEB: all tasks are Retrieval type, so Retrieval column = Mean (Task)
-    if "Retrieval" in joint_table.columns:
-        rename_dict["Retrieval"] = "Mean (Task)"
+
     joint_table = joint_table.rename(columns=rename_dict)
 
     # Move borda rank to front

mteb/benchmarks/benchmark.py

@@ -87,7 +87,10 @@ class RtebBenchmark(Benchmark):
     def _create_summary_table(
         self, benchmark_results: BenchmarkResults
     ) -> pd.DataFrame:
-        return _create_summary_table_mean_public_private(benchmark_results)
+        joint_table = _create_summary_table_mean_public_private(benchmark_results)
+        # For RTEB: all tasks are Retrieval type, so Retrieval column = Mean (Task)
+        joint_table = joint_table.rename(columns={"Retrieval": "Mean (Task)"})
+        return joint_table
 
 
 class HUMEBenchmark(Benchmark):
@@ -106,3 +109,17 @@ class MIEBBenchmark(Benchmark):
         self, benchmark_results: BenchmarkResults
     ) -> pd.DataFrame:
         return _create_summary_table_mean_task_type(benchmark_results)
+
+
+class VidoreBenchmark(Benchmark):
+    """Wrapper for Vidore3 benchmark."""
+
+    def _create_summary_table(
+        self, benchmark_results: BenchmarkResults
+    ) -> pd.DataFrame:
+        joint_table = _create_summary_table_mean_public_private(benchmark_results)
+        # For ViDoRe (V1, V2, V3): all tasks are Document Understanding type, so Document Understanding column = Mean (Task)
+        joint_table = joint_table.rename(
+            columns={"Document Understanding": "Mean (Task)"}
+        )
+        return joint_table

mteb/benchmarks/benchmarks/__init__.py

@@ -27,6 +27,7 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     MTEB_KOR,
     MTEB_MAIN_RU,
     MTEB_MINERS_BITEXT_MINING,
+    MTEB_NL,
     MTEB_POL,
     MTEB_RETRIEVAL_LAW,
     MTEB_RETRIEVAL_MEDICAL,
@@ -37,6 +38,7 @@ from mteb.benchmarks.benchmarks.benchmarks import (
     SEB,
     VIDORE,
     VIDORE_V2,
+    VIDORE_V3,
     VISUAL_DOCUMENT_RETRIEVAL,
     VN_MTEB,
     CoIR,
@@ -87,6 +89,7 @@ __all__ = [
     "MTEB_KOR",
     "MTEB_MAIN_RU",
     "MTEB_MINERS_BITEXT_MINING",
+    "MTEB_NL",
     "MTEB_POL",
     "MTEB_RETRIEVAL_LAW",
     "MTEB_RETRIEVAL_MEDICAL",
@@ -106,6 +109,7 @@ __all__ = [
     "SEB",
     "VIDORE",
     "VIDORE_V2",
+    "VIDORE_V3",
     "VISUAL_DOCUMENT_RETRIEVAL",
     "VN_MTEB",
     "CoIR",