mteb 2.7.2__py3-none-any.whl → 2.7.9__py3-none-any.whl

mteb/_evaluators/sklearn_evaluator.py

@@ -1,18 +1,22 @@
-import logging
-from typing import Any, Protocol, cast
+from __future__ import annotations
 
-import numpy as np
-from datasets import Dataset
-from torch.utils.data import DataLoader
-from typing_extensions import Self
+import logging
+from typing import TYPE_CHECKING, Any, Protocol, cast
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import Array, BatchedInput, EncodeKwargs
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    import numpy as np
+    from datasets import Dataset
+    from torch.utils.data import DataLoader
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, BatchedInput, EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 
@@ -50,18 +54,20 @@ class SklearnEvaluator(Evaluator):
         self.evaluator_model = evaluator_model
 
     def create_dataloaders(
-        self, encode_kwargs: EncodeKwargs
+        self, encode_kwargs: EncodeKwargs, num_proc: int
     ) -> tuple[DataLoader[BatchedInput], DataLoader[BatchedInput]]:
         dataloader_train = create_dataloader(
             self.train_dataset,
             self.task_metadata,
             input_column=self.values_column_name,
+            num_proc=num_proc,
             **encode_kwargs,
         )
         dataloader_test = create_dataloader(
             self.eval_dataset,
             self.task_metadata,
             input_column=self.values_column_name,
+            num_proc=num_proc,
             **encode_kwargs,
         )
         return dataloader_train, dataloader_test
@@ -72,6 +78,7 @@ class SklearnEvaluator(Evaluator):
         *,
         encode_kwargs: EncodeKwargs,
         test_cache: Array | None = None,
+        num_proc: int = 1,
     ) -> tuple[np.ndarray, Array]:
         """Classification evaluation by training a sklearn classifier on the embeddings of the training set and evaluating on the embeddings of the test set.
 
@@ -79,6 +86,7 @@ class SklearnEvaluator(Evaluator):
             model: Encoder
             encode_kwargs: encode kwargs
             test_cache: embeddings of the test set, if already computed
+            num_proc: number of processes to use
 
         Returns:
             Tuple of test predictions and embeddings
@@ -86,6 +94,7 @@ class SklearnEvaluator(Evaluator):
         """
         dataloader_train, dataloader_test = self.create_dataloaders(
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
 
         logger.info("Running - Encoding samples...")
@@ -104,7 +113,7 @@ class SklearnEvaluator(Evaluator):
                 hf_subset=self.hf_subset,
                 **encode_kwargs,
             )
-            test_cache = cast(Array, test_cache)
+            test_cache = cast("Array", test_cache)
 
         logger.info("Running - Fitting classifier...")
         y_train = self.train_dataset[self.label_column_name]

mteb/_evaluators/text/bitext_mining_evaluator.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import torch
 from datasets import Dataset
@@ -6,9 +9,11 @@ from tqdm.auto import tqdm
 
 from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import Array, EncodeKwargs
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, EncodeKwargs
 
 logger = logging.getLogger(__name__)
 
@@ -36,6 +41,7 @@ class BitextMiningEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> dict[str, list[dict[str, float]]]:
         pair_elements = {p for pair in self.pairs for p in pair}
         if isinstance(self.sentences, Dataset):
@@ -50,6 +56,7 @@ class BitextMiningEvaluator(Evaluator):
         for sub in tqdm(subsets):
             dataloader = _create_dataloader_from_texts(
                 self.sentences[sub],
+                num_proc=num_proc,
                 **encode_kwargs,
             )
             embeddings[sub] = model.encode(

mteb/_evaluators/text/summarization_evaluator.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import logging
 import sys
-from typing import TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 import numpy as np
 import torch
@@ -9,10 +11,12 @@ from tqdm.auto import tqdm
 
 from mteb._create_dataloaders import _create_dataloader_from_texts
 from mteb._evaluators.evaluator import Evaluator
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import cos_sim, dot_score
-from mteb.types import EncodeKwargs
+
+if TYPE_CHECKING:
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs
 
 # if later than python 3.13 use typing module
 if sys.version_info >= (3, 13):
@@ -96,6 +100,7 @@ class SummarizationEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> SummarizationDistances:
         # Get the human & machine summaries for the text in one go for all
         human_lens = [len(human_summaries) for human_summaries in self.human_summaries]
@@ -111,6 +116,7 @@ class SummarizationEvaluator(Evaluator):
                     for human_summaries in self.human_summaries
                     for summary in human_summaries
                 ],
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 from datasets import Dataset
 
@@ -6,13 +9,17 @@ from mteb._create_dataloaders import (
     _create_dataloader_from_texts,
     create_dataloader,
 )
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
 from mteb.similarity_functions import similarity
-from mteb.types import Array, EncodeKwargs
 
 from .evaluator import Evaluator
 
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import Array, EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 
@@ -41,11 +48,13 @@ class ZeroShotClassificationEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> Array:
         dataloader = create_dataloader(
             self.dataset,
             input_column=self.input_column_name,
             task_metadata=self.task_metadata,
+            num_proc=num_proc,
             **encode_kwargs,
         )
 

mteb/_helpful_enum.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
 from enum import Enum
+from typing import TYPE_CHECKING
 
-from typing_extensions import Self
+if TYPE_CHECKING:
+    from typing_extensions import Self
 
 
 class HelpfulStrEnum(str, Enum):

mteb/abstasks/_data_filter/filters.py

@@ -1,12 +1,18 @@
 """Simplified version of https://gist.github.com/AlexeyVatolin/ea3adc21aa7a767603ff393b22085adc from https://github.com/embeddings-benchmark/mteb/pull/2900"""
 
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 import datasets
 import pandas as pd
-from datasets import Dataset, DatasetDict
+from datasets import DatasetDict
+
+if TYPE_CHECKING:
+    from datasets import Dataset
 
-from mteb import TaskMetadata
+    from mteb import TaskMetadata
 
 logger = logging.getLogger(__name__)
 

mteb/abstasks/_data_filter/task_pipelines.py

@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
 from datasets import DatasetDict
 
-from mteb import TaskMetadata
-from mteb.abstasks import AbsTaskClassification
 from mteb.abstasks._data_filter.filters import (
     deduplicate,
     filter_empty,
@@ -13,6 +14,10 @@ from mteb.abstasks._data_filter.filters import (
     split_train_test,
 )
 
+if TYPE_CHECKING:
+    from mteb import TaskMetadata
+    from mteb.abstasks import AbsTaskClassification
+
 logger = logging.getLogger(__name__)
 
 

mteb/abstasks/_statistics_calculation.py

@@ -2,10 +2,8 @@ from __future__ import annotations
 
 import hashlib
 from collections import Counter
-from collections.abc import Mapping
 from typing import TYPE_CHECKING, cast
 
-from mteb.types import TopRankedDocumentsType
 from mteb.types.statistics import (
     ImageStatistics,
     LabelStatistics,
@@ -16,8 +14,12 @@ from mteb.types.statistics import (
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
+
     from PIL import Image
 
+    from mteb.types import TopRankedDocumentsType
+
 
 def calculate_text_statistics(texts: list[str]) -> TextStatistics:
     """Calculate descriptive statistics for a list of texts.
@@ -87,13 +89,13 @@ def calculate_label_statistics(labels: list[int | list[int]]) -> LabelStatistics
 
     if not isinstance(labels[0], list):
         # single label classification
-        single_label = cast(list[int], labels)
+        single_label = cast("list[int]", labels)
         label_len = [1] * len(single_label)
         total_label_len = len(single_label)
         total_labels.extend(single_label)
     elif isinstance(labels[0], list):
         # multilabel classification
-        multilabel_labels = cast(list[list[int]], labels)
+        multilabel_labels = cast("list[list[int]]", labels)
         label_len = [len(l) for l in multilabel_labels]
         total_label_len = sum(label_len)
         for l in multilabel_labels:

mteb/abstasks/abstask.py

@@ -1,30 +1,38 @@
+from __future__ import annotations
+
 import json
 import logging
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 from copy import copy
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
 from datasets import ClassLabel, Dataset, DatasetDict, load_dataset
 from sklearn.preprocessing import MultiLabelBinarizer
 from tqdm.auto import tqdm
-from typing_extensions import Self
 
 from mteb._set_seed import _set_seed
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.languages import LanguageScripts
 from mteb.models import (
     CrossEncoderProtocol,
     EncoderProtocol,
-    MTEBModels,
     SearchProtocol,
 )
-from mteb.types import HFSubset, Modalities, ScoresDict
-from mteb.types._encoder_io import EncodeKwargs
-from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from typing_extensions import Self
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import (
+        MTEBModels,
+    )
+    from mteb.types import EncodeKwargs, HFSubset, Modalities, ScoresDict
+    from mteb.types.statistics import DescriptiveStatistics, SplitDescriptiveStatistics
 
 logger = logging.getLogger(__name__)
 
@@ -108,11 +116,14 @@ class AbsTask(ABC):
             logger.warning(msg)
             warnings.warn(msg)
 
-    def dataset_transform(self):
+    def dataset_transform(self, num_proc: int = 1):
         """A transform operations applied to the dataset after loading.
 
         This method is useful when the dataset from Huggingface is not in an `mteb` compatible format.
         Override this method if your dataset requires additional transformation.
+
+        Args:
+            num_proc: Number of processes to use for the transformation.
         """
         pass
 
@@ -124,6 +135,7 @@ class AbsTask(ABC):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluates an MTEB compatible model on the task.
@@ -134,6 +146,7 @@ class AbsTask(ABC):
             subsets_to_run: List of huggingface subsets (HFSubsets) to evaluate. If None, all subsets are evaluated.
             encode_kwargs: Additional keyword arguments that are passed to the model's `encode` method.
             prediction_folder: Folder to save model predictions
+            num_proc: Number of processes to use for loading the dataset or processing.
             kwargs: Additional keyword arguments that are passed to the _evaluate_subset method.
 
         Returns:
@@ -163,7 +176,7 @@ class AbsTask(ABC):
         if not self.data_loaded:
             self.load_data()
 
-        self.dataset = cast(dict[HFSubset, DatasetDict], self.dataset)
+        self.dataset = cast("dict[HFSubset, DatasetDict]", self.dataset)
 
         scores = {}
         if self.hf_subsets is None:
@@ -189,6 +202,7 @@ class AbsTask(ABC):
                 hf_subset=hf_subset,
                 encode_kwargs=encode_kwargs,
                 prediction_folder=prediction_folder,
+                num_proc=num_proc,
                 **kwargs,
             )
             self._add_main_score(scores[hf_subset])
@@ -204,6 +218,7 @@ class AbsTask(ABC):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
+        num_proc: int = 1,
         **kwargs: Any,
     ) -> ScoresDict:
         raise NotImplementedError(
@@ -308,11 +323,15 @@ class AbsTask(ABC):
             )  # only take the specified test split.
         return dataset_dict
 
-    def load_data(self) -> None:
+    def load_data(self, num_proc: int = 1, **kwargs: Any) -> None:
         """Loads dataset from HuggingFace hub
 
         This is the main loading function for Task. Do not overwrite this, instead we recommend using `dataset_transform`, which is called after the
         dataset is loaded using `datasets.load_dataset`.
+
+        Args:
+            num_proc: Number of processes to use for loading the dataset.
+            kwargs: Additional keyword arguments passed to the load_dataset function. Keep for forward compatibility.
         """
         if self.data_loaded:
             return
@@ -325,11 +344,12 @@ class AbsTask(ABC):
                     self.dataset[hf_subset] = load_dataset(
                         name=hf_subset,
                         **self.metadata.dataset,
+                        num_proc=num_proc,
                     )
         else:
             # some of monolingual datasets explicitly adding the split name to the dataset name
-            self.dataset = load_dataset(**self.metadata.dataset)
-        self.dataset_transform()
+            self.dataset = load_dataset(**self.metadata.dataset, num_proc=num_proc)
+        self.dataset_transform(num_proc=num_proc)
         self.data_loaded = True
 
     def fast_load(self) -> None:
@@ -352,12 +372,13 @@ class AbsTask(ABC):
             self.dataset[lang] = DatasetDict(subset)
 
     def calculate_descriptive_statistics(
-        self, overwrite_results: bool = False
+        self, overwrite_results: bool = False, num_proc: int = 1
     ) -> dict[str, DescriptiveStatistics]:
         """Calculates descriptive statistics from the dataset.
 
         Args:
             overwrite_results: Whether to overwrite existing results. If False and results already exist, the existing results will be loaded from cache.
+            num_proc: Number of processes to use for loading the dataset.
 
         Returns:
             A dictionary containing descriptive statistics for each split.
@@ -371,7 +392,7 @@ class AbsTask(ABC):
             return existing_stats
 
         if not self.data_loaded:
-            self.load_data()
+            self.load_data(num_proc=num_proc)
 
         descriptive_stats: dict[str, DescriptiveStatistics] = {}
         hf_subset_stat: Literal["hf_subset_descriptive_stats"] = (
@@ -509,7 +530,7 @@ class AbsTask(ABC):
         scores["main_score"] = scores[self.metadata.main_score]
 
     def _upload_dataset_to_hub(
-        self, repo_name: str, fields: list[str] | dict[str, str]
+        self, repo_name: str, fields: list[str] | dict[str, str], num_proc: int = 1
     ) -> None:
         if self.dataset is None:
             raise ValueError("Dataset not loaded")
@@ -534,7 +555,10 @@ class AbsTask(ABC):
                         )
                 sentences = DatasetDict(sentences)
                 sentences.push_to_hub(
-                    repo_name, config, commit_message=f"Add {config} dataset"
+                    repo_name,
+                    config,
+                    commit_message=f"Add {config} dataset",
+                    num_proc=num_proc,
                 )
         else:
             sentences = {}
@@ -551,16 +575,19 @@ class AbsTask(ABC):
                         {field: self.dataset[split][field] for field in fields}
                     )
             sentences = DatasetDict(sentences)
-            sentences.push_to_hub(repo_name, commit_message="Add dataset")
+            sentences.push_to_hub(
+                repo_name, commit_message="Add dataset", num_proc=num_proc
+            )
 
-    def _push_dataset_to_hub(self, repo_name: str) -> None:
+    def _push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         raise NotImplementedError
 
-    def push_dataset_to_hub(self, repo_name: str) -> None:
+    def push_dataset_to_hub(self, repo_name: str, num_proc: int = 1) -> None:
         """Push the dataset to the HuggingFace Hub.
 
         Args:
             repo_name: The name of the repository to push the dataset to.
+            num_proc: Number of processes to use for loading the dataset.
 
         Examples:
             >>> import mteb
@@ -572,7 +599,7 @@ class AbsTask(ABC):
         if not self.data_loaded:
             self.load_data()
 
-        self._push_dataset_to_hub(repo_name)
+        self._push_dataset_to_hub(repo_name, num_proc)
         # dataset repo not creating when pushing card
         self.metadata.push_dataset_card_to_hub(repo_name)
 

mteb/abstasks/aggregate_task_metadata.py

@@ -1,28 +1,39 @@
+from __future__ import annotations
+
 import logging
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 from pydantic import ConfigDict, Field, model_validator
-from typing_extensions import Self
 
 from mteb.types import (
-    ISOLanguageScript,
     Languages,
-    Licenses,
-    Modalities,
-    StrDate,
 )
 
 from .abstask import AbsTask
 from .task_metadata import (
-    AnnotatorType,
     MetadataDatasetDict,
-    SampleCreationMethod,
-    TaskDomain,
     TaskMetadata,
-    TaskSubtype,
     TaskType,
 )
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from mteb.types import (
+        ISOLanguageScript,
+        Licenses,
+        Modalities,
+        StrDate,
+    )
+
+    from .task_metadata import (
+        AnnotatorType,
+        SampleCreationMethod,
+        TaskDomain,
+        TaskSubtype,
+    )
+
 logger = logging.getLogger(__name__)
 
 

mteb/abstasks/aggregated_task.py

@@ -1,19 +1,26 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Mapping
-from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
-from datasets import Dataset, DatasetDict
 
-from mteb.models.models_protocols import MTEBModels
 from mteb.results.task_result import TaskResult
-from mteb.types import EncodeKwargs, HFSubset, ScoresDict
-from mteb.types.statistics import DescriptiveStatistics
 
 from .abstask import AbsTask
-from .aggregate_task_metadata import AggregateTaskMetadata
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from pathlib import Path
+
+    from datasets import Dataset, DatasetDict
+
+    from mteb.models.models_protocols import MTEBModels
+    from mteb.types import EncodeKwargs, HFSubset, ScoresDict
+    from mteb.types.statistics import DescriptiveStatistics
+
+    from .aggregate_task_metadata import AggregateTaskMetadata
 
 logger = logging.getLogger(__name__)