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

mteb/_create_dataloaders.py

@@ -1,21 +1,28 @@
+from __future__ import annotations
+
 import logging
 import warnings
-from collections.abc import Callable
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from datasets import Dataset, Image
 from torch.utils.data import DataLoader, default_collate
 
-from mteb.abstasks.task_metadata import TaskMetadata
 from mteb.types import (
-    BatchedInput,
-    Conversation,
     ConversationTurn,
     PromptType,
-    QueryDatasetType,
 )
-from mteb.types._encoder_io import CorpusInput, ImageInput, QueryInput, TextInput
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.types import (
+        BatchedInput,
+        Conversation,
+        QueryDatasetType,
+    )
+    from mteb.types._encoder_io import CorpusInput, ImageInput, QueryInput, TextInput
 
 logger = logging.getLogger(__name__)
 
@@ -23,6 +30,7 @@ logger = logging.getLogger(__name__)
 def _create_dataloader_from_texts(
     text: list[str],
     batch_size: int = 32,
+    num_proc: int = 1,
     **kwargs: Any,
 ) -> DataLoader[TextInput]:
     """Create a dataloader from a list of text.
@@ -30,15 +38,17 @@ def _create_dataloader_from_texts(
     Args:
         text: A list of text to create a dataloader from.
         batch_size: Batch size for the dataloader.
+        num_proc: Number of processes to use.
         kwargs: Not used, present catching extra arguments.
 
     Returns:
         A dataloader with the text.
     """
     dataset = Dataset.from_dict({"text": text})
-    return torch.utils.data.DataLoader(
+    return DataLoader(
         dataset,
         batch_size=batch_size,
+        num_workers=num_proc if num_proc > 1 else 0,
     )
 
 
@@ -64,20 +74,27 @@ def _corpus_to_dict(
 def _create_dataloader_for_retrieval_corpus(
     dataset: Dataset,
     batch_size: int = 32,
+    num_proc: int = 1,
 ) -> DataLoader[CorpusInput]:
     """Create a dataloader from a corpus.
 
     Args:
         dataset: Corpus
         batch_size: Batch size for the dataloader.
+        num_proc: Number of processes to use.
 
     Returns:
         A dataloader with the corpus.
     """
-    new_ds = dataset.map(_corpus_to_dict, desc="Converting corpus dict")
-    return torch.utils.data.DataLoader(
+    new_ds = dataset.map(
+        _corpus_to_dict,
+        desc="Converting corpus dict",
+        num_proc=num_proc,
+    )
+    return DataLoader(
         new_ds,
         batch_size=batch_size,
+        num_workers=num_proc if num_proc > 1 else 0,
     )
 
 
@@ -94,12 +111,14 @@ def _combine_queries_with_instruction_text(row: dict[str, str]) -> dict[str, str
 def _create_text_dataloader_for_queries(
     queries: QueryDatasetType,
     batch_size: int = 32,
+    num_proc: int = 1,
 ) -> DataLoader[QueryInput]:
     """Create a dataloader from a list of queries.
 
     Args:
         queries: A list of queries.
         batch_size: Batch size for the dataloader.
+        num_proc: Number of processes to use.
 
     Returns:
         A dataloader with the queries.
@@ -107,10 +126,12 @@ def _create_text_dataloader_for_queries(
     queries = queries.map(
         _combine_queries_with_instruction_text,
         desc="Processing queries for dataloading",
+        num_proc=num_proc,
     )
-    return torch.utils.data.DataLoader(
+    return DataLoader(
         queries,
         batch_size=batch_size,
+        num_workers=num_proc if num_proc > 1 else 0,
     )
 
 
@@ -128,7 +149,7 @@ def _convert_conv_history_to_query(
     conversation = row["text"]
     # if it's a list of strings, just join them
     if isinstance(conversation, list) and isinstance(conversation[0], str):
-        conversation_ = cast(list[str], conversation)
+        conversation_ = cast("list[str]", conversation)
         conv_str = "; ".join(conversation_)
         current_conversation = [
             ConversationTurn(role="user", content=message) for message in conversation_
@@ -173,18 +194,20 @@ def _convert_conv_history_to_query(
 
     row["text"] = conv_str
     row["conversation"] = current_conversation
-    return cast(dict[str, str | list[ConversationTurn]], row)
+    return cast("dict[str, str | list[ConversationTurn]]", row)
 
 
 def _create_dataloader_for_queries_conversation(
     queries: QueryDatasetType,
     batch_size: int = 32,
+    num_proc: int = 1,
 ) -> DataLoader[QueryInput]:
     """Create a dataloader from a list of queries.
 
     Args:
         queries: A list of queries.
         batch_size: Batch size for the dataloader.
+        num_proc: Number of processes to use.
 
     Returns:
         A dataloader with the queries.
@@ -193,9 +216,11 @@ def _create_dataloader_for_queries_conversation(
         queries.map(
             _convert_conv_history_to_query,
             desc="Converting conversations to queries",
+            num_proc=num_proc,
         ),
         collate_fn=_custom_collate_fn,
         batch_size=batch_size,
+        num_workers=num_proc if num_proc > 1 else 0,
     )
 
 
@@ -240,6 +265,7 @@ def _prepare_image_dataset(
     dataset: Dataset,
     image_column_name: str | None = None,
     transform: Callable[[Any], Any] | None = None,
+    num_proc: int = 1,
 ) -> Dataset:
     """Prepare the image dataset by converting images to RGB and applying transformations."""
     if (
@@ -255,6 +281,7 @@ def _prepare_image_dataset(
         _convert_images_to_rgb,
         fn_kwargs={"image_col_name": "image", "transform": transform},
         desc="Converting images to RGB",
+        num_proc=num_proc,
     )
 
 
@@ -288,6 +315,7 @@ def _create_image_dataloader(
     batch_size: int = 32,
     transform: Callable[[Any], Any] | None = None,
     collate_fn: Callable[[list[dict[str, Any]]], dict[str, Any]] = _custom_collate_fn,
+    num_proc: int = 1,
 ) -> DataLoader[ImageInput]:
     """Creates a DataLoader with the image dataset prepared using the explicit transformation.
 
@@ -297,33 +325,41 @@ def _create_image_dataloader(
         batch_size: Batch size for the dataloader.
         transform: A transformation function to apply to each image (e.g., converting to tensor).
         collate_fn: A custom collate function to handle batching.
+        num_proc: Number of processes to use.
 
     Returns:
         A DataLoader with the image dataset.
     """
     dataset = _prepare_image_dataset(
-        dataset, image_column_name, transform
+        dataset,
+        image_column_name,
+        transform,
+        num_proc=num_proc,
     ).select_columns(["image"])
     return DataLoader(
         dataset,
         batch_size=batch_size,
         collate_fn=collate_fn,
         shuffle=False,
+        num_workers=num_proc if num_proc > 1 else 0,
     )
 
 
 def _create_text_queries_dataloader(
     dataset: Dataset,
     batch_size: int = 32,
+    num_proc: int = 1,
 ) -> DataLoader[QueryInput]:
     if not isinstance(dataset["text"][0], list):
         return _create_text_dataloader_for_queries(
             dataset,
             batch_size=batch_size,
+            num_proc=num_proc,
         )
     return _create_dataloader_for_queries_conversation(
         dataset,
         batch_size=batch_size,
+        num_proc=num_proc,
     )
 
 
@@ -332,6 +368,7 @@ def _create_queries_dataloader(
     task_metadata: TaskMetadata,
     input_column: str | None = None,
     batch_size: int = 32,
+    num_proc: int = 1,
 ) -> DataLoader[QueryInput | ImageInput]:
     """Create a dataloader for queries."""
     queries_type = task_metadata.get_modalities(PromptType.query)
@@ -339,12 +376,14 @@ def _create_queries_dataloader(
         return _create_text_queries_dataloader(
             dataset,
             batch_size=batch_size,
+            num_proc=num_proc,
         )
     if "image" in queries_type:  # contains image
         return _create_image_dataloader(
             dataset,
             image_column_name="image",
             batch_size=batch_size,
+            num_proc=num_proc,
         )
     raise ValueError(f"Can't handle queries type {queries_type}")
 
@@ -354,6 +393,7 @@ def _create_document_dataloader(
     task_metadata: TaskMetadata,
     input_column: str | None = None,
     batch_size: int = 32,
+    num_proc: int = 1,
 ) -> DataLoader[CorpusInput | ImageInput]:
     """Create a dataloader for documents.
 
@@ -362,6 +402,7 @@ def _create_document_dataloader(
         task_metadata: Metadata of the task to determine the document type.
         input_column: The column to use as input. If None, it will use the first column that matches the modality.
         batch_size: Batch size for the dataloader.
+        num_proc: Number of processes to use.
 
     Returns:
         A dataloader for the documents.
@@ -371,12 +412,14 @@ def _create_document_dataloader(
         return _create_dataloader_for_retrieval_corpus(
             dataset,
             batch_size=batch_size,
+            num_proc=num_proc,
         )
     if "image" in document_type:  # contains image
         return _create_image_dataloader(
             dataset,
             image_column_name="image",
             batch_size=batch_size,
+            num_proc=num_proc,
         )
     raise ValueError(f"Can't handle queries type {document_type}")
 
@@ -387,6 +430,7 @@ def create_dataloader(
     prompt_type: PromptType | None = None,
     input_column: str | None = None,
     batch_size: int = 32,
+    num_proc: int = 1,
     **kwargs: Any,
 ) -> DataLoader[BatchedInput]:
     """Create a dataloader from a dataset.
@@ -400,6 +444,7 @@ def create_dataloader(
         prompt_type: The type of prompt to create a dataloader for. If None, it will be inferred from the task metadata.
         input_column: The column to use as input. If None, it will use the first column that matches the modality.
         batch_size: The batch size for the dataloader.
+        num_proc: The number of processes to use for dataset processing.
         **kwargs: Additional arguments to pass to the dataloader creation functions.
 
     Returns:
@@ -411,6 +456,7 @@ def create_dataloader(
             task_metadata,
             batch_size=batch_size,
             input_column=input_column,
+            num_proc=num_proc,
         )
     if prompt_type == PromptType.document:
         return _create_document_dataloader(
@@ -418,6 +464,7 @@ def create_dataloader(
             task_metadata,
             input_column=input_column,
             batch_size=batch_size,
+            num_proc=num_proc,
         )
 
     if "image" in task_metadata.modalities:
@@ -425,6 +472,7 @@ def create_dataloader(
             dataset,
             image_column_name=input_column,
             batch_size=batch_size,
+            num_proc=num_proc,
         )
     if "text" in task_metadata.modalities and input_column is not None:
         return _create_dataloader_from_texts(
@@ -434,4 +482,5 @@ def create_dataloader(
     return DataLoader(
         dataset,
         batch_size=batch_size,
+        num_workers=num_proc if num_proc > 1 else 0,
     )

mteb/_evaluators/any_sts_evaluator.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
-from typing import TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
-from datasets import Dataset
 from sklearn.metrics.pairwise import (
     paired_cosine_distances,
     paired_euclidean_distances,
@@ -9,13 +10,17 @@ from sklearn.metrics.pairwise import (
 )
 
 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 EncodeKwargs, PromptType
 
 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 EncodeKwargs, PromptType
+
 logger = logging.getLogger(__name__)
 
 
@@ -61,6 +66,7 @@ class AnySTSEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> STSEvaluatorScores:
         logger.info("Running semantic similarity - Encoding samples (1/2)")
         embeddings1 = model.encode(
@@ -68,6 +74,7 @@ class AnySTSEvaluator(Evaluator):
                 self.dataset,
                 self.task_metadata,
                 input_column=self.input_columns[0],
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,

mteb/_evaluators/clustering_evaluator.py

@@ -1,15 +1,21 @@
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
 
-from datasets import Dataset
 from sklearn import cluster
 
 from mteb._create_dataloaders import create_dataloader
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import EncoderProtocol
-from mteb.types import 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 EncodeKwargs
+
 logger = logging.getLogger(__name__)
 
 
@@ -39,11 +45,13 @@ class ClusteringEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> list[int]:
         data_loader = create_dataloader(
             self.dataset,
             self.task_metadata,
             input_column=self.input_column_name,
+            num_proc=num_proc,
             **encode_kwargs,
         )
 

mteb/_evaluators/evaluator.py

@@ -1,10 +1,15 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from mteb.abstasks.abstask import _set_seed
-from mteb.models import EncoderProtocol
-from mteb.types import EncodeKwargs
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Mapping
+
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs
 
 
 class Evaluator(ABC):
@@ -19,7 +24,7 @@ class Evaluator(ABC):
 
     @abstractmethod
     def __call__(
-        self, model: EncoderProtocol, *, encode_kwargs: EncodeKwargs
+        self, model: EncoderProtocol, *, encode_kwargs: EncodeKwargs, num_proc: int = 1
     ) -> Mapping[str, float] | Iterable[Any]:
         """This is called during training to evaluate the model.
 
@@ -28,5 +33,6 @@ class Evaluator(ABC):
         Args:
             model: the model to evaluate
             encode_kwargs: kwargs to pass to the model's encode method
+            num_proc: number of processes to use for data loading
         """
         pass

mteb/_evaluators/image/imagetext_pairclassification_evaluator.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -14,13 +13,16 @@ from mteb._create_dataloaders import (
 )
 from mteb._evaluators.evaluator import Evaluator
 from mteb._requires_package import requires_image_dependencies
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models.models_protocols import EncoderProtocol
-from mteb.types import EncodeKwargs
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from PIL.Image import Image
 
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models.models_protocols import EncoderProtocol
+    from mteb.types import EncodeKwargs
+
 
 logger = logging.getLogger(__name__)
 
@@ -89,6 +91,7 @@ class ImageTextPairClassificationEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> list[torch.Tensor]:
         images = []
         if isinstance(self.images_column_names, str):
@@ -111,6 +114,7 @@ class ImageTextPairClassificationEvaluator(Evaluator):
         text_embeddings = model.encode(
             _create_dataloader_from_texts(
                 texts,
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
@@ -127,10 +131,15 @@ class ImageTextPairClassificationEvaluator(Evaluator):
             dim=-1,
         ).view(len(self.dataset), self.num_texts_per_sample, -1)
 
+        def _image_collate_fn(batch):
+            """Collate function for image batches."""
+            return {"image": [item["image"] for item in batch]}
+
         image_embeddings = model.encode(
             DataLoader(
                 CustomImageDataset(images),
-                collate_fn=lambda x: {"image": [item["image"] for item in x]},
+                collate_fn=_image_collate_fn,
+                num_workers=num_proc if num_proc > 1 else 0,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,

mteb/_evaluators/pair_classification_evaluator.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
+
 import logging
-from typing import Any, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 
 import numpy as np
-from datasets import Dataset
 from sklearn.metrics.pairwise import (
     paired_cosine_distances,
     paired_euclidean_distances,
@@ -11,10 +12,14 @@ from sklearn.metrics.pairwise import (
 
 from mteb._create_dataloaders import _create_dataloader_from_texts, create_dataloader
 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 EncodeKwargs, PromptType
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import EncoderProtocol
+    from mteb.types import EncodeKwargs, PromptType
 
 logger = logging.getLogger(__name__)
 
@@ -86,6 +91,7 @@ class PairClassificationEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> PairClassificationDistances:
         logger.info("Running pair classification - Encoding samples (1/2)")
         embeddings1 = model.encode(
@@ -93,6 +99,7 @@ class PairClassificationEvaluator(Evaluator):
                 self.dataset,
                 task_metadata=self.task_metadata,
                 input_column=self.input1_column_name,
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,
@@ -107,6 +114,7 @@ class PairClassificationEvaluator(Evaluator):
                 self.dataset,
                 task_metadata=self.task_metadata,
                 input_column=self.input2_column_name,
+                num_proc=num_proc,
                 **encode_kwargs,
             ),
             task_metadata=self.task_metadata,

mteb/_evaluators/retrieval_evaluator.py

@@ -1,23 +1,29 @@
-import logging
-from collections.abc import Sequence
+from __future__ import annotations
 
-from mteb.abstasks.task_metadata import TaskMetadata
-from mteb.models import SearchProtocol
-from mteb.types import (
-    CorpusDatasetType,
-    EncodeKwargs,
-    QueryDatasetType,
-    RelevantDocumentsType,
-    RetrievalEvaluationResult,
-    RetrievalOutputType,
-    TopRankedDocumentsType,
-)
+import logging
+from typing import TYPE_CHECKING
 
 from .evaluator import Evaluator
 from .retrieval_metrics import (
     calculate_retrieval_scores,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from mteb.abstasks.task_metadata import TaskMetadata
+    from mteb.models import SearchProtocol
+    from mteb.types import (
+        CorpusDatasetType,
+        EncodeKwargs,
+        QueryDatasetType,
+        RelevantDocumentsType,
+        RetrievalEvaluationResult,
+        RetrievalOutputType,
+        TopRankedDocumentsType,
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -49,6 +55,7 @@ class RetrievalEvaluator(Evaluator):
         self,
         search_model: SearchProtocol,
         encode_kwargs: EncodeKwargs,
+        num_proc: int = 1,
     ) -> RetrievalOutputType:
         logger.info("Running retrieval task - Indexing corpus...")
         search_model.index(
@@ -57,6 +64,7 @@ class RetrievalEvaluator(Evaluator):
             hf_split=self.hf_split,
             hf_subset=self.hf_subset,
             encode_kwargs=encode_kwargs,
+            num_proc=num_proc,
         )
         logger.info("Running retrieval task - Searching queries...")
         return search_model.search(
@@ -67,6 +75,7 @@ class RetrievalEvaluator(Evaluator):
             hf_subset=self.hf_subset,
             encode_kwargs=encode_kwargs,
             top_ranked=self.top_ranked,
+            num_proc=num_proc,
         )
 
     def evaluate(

mteb/_evaluators/retrieval_metrics.py

@@ -1,7 +1,8 @@
+from __future__ import annotations
+
 import logging
 from collections import defaultdict
-from collections.abc import Mapping
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import pandas as pd
@@ -9,7 +10,12 @@ import pytrec_eval
 from packaging.version import Version
 from sklearn.metrics import auc
 
-from mteb.types import RelevantDocumentsType, RetrievalEvaluationResult
+from mteb.types import RetrievalEvaluationResult
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from mteb.types import RelevantDocumentsType
 
 logger = logging.getLogger(__name__)