mteb 2.7.17__py3-none-any.whl → 2.7.19__py3-none-any.whl

mteb/_create_dataloaders.py

@@ -30,7 +30,7 @@ logger = logging.getLogger(__name__)
 def _create_dataloader_from_texts(
     text: list[str],
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
     **kwargs: Any,
 ) -> DataLoader[TextInput]:
     """Create a dataloader from a list of text.
@@ -48,7 +48,7 @@ def _create_dataloader_from_texts(
     return DataLoader(
         dataset,
         batch_size=batch_size,
-        num_workers=num_proc if num_proc > 1 else 0,
+        num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
     )
 
 
@@ -74,7 +74,7 @@ def _corpus_to_dict(
 def _create_dataloader_for_retrieval_corpus(
     dataset: Dataset,
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> DataLoader[CorpusInput]:
     """Create a dataloader from a corpus.
 
@@ -94,7 +94,7 @@ def _create_dataloader_for_retrieval_corpus(
     return DataLoader(
         new_ds,
         batch_size=batch_size,
-        num_workers=num_proc if num_proc > 1 else 0,
+        num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
     )
 
 
@@ -111,7 +111,7 @@ 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,
+    num_proc: int | None = None,
 ) -> DataLoader[QueryInput]:
     """Create a dataloader from a list of queries.
 
@@ -131,7 +131,7 @@ def _create_text_dataloader_for_queries(
     return DataLoader(
         queries,
         batch_size=batch_size,
-        num_workers=num_proc if num_proc > 1 else 0,
+        num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
     )
 
 
@@ -200,7 +200,7 @@ def _convert_conv_history_to_query(
 def _create_dataloader_for_queries_conversation(
     queries: QueryDatasetType,
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> DataLoader[QueryInput]:
     """Create a dataloader from a list of queries.
 
@@ -220,7 +220,7 @@ def _create_dataloader_for_queries_conversation(
         ),
         collate_fn=_custom_collate_fn,
         batch_size=batch_size,
-        num_workers=num_proc if num_proc > 1 else 0,
+        num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
     )
 
 
@@ -265,7 +265,7 @@ def _prepare_image_dataset(
     dataset: Dataset,
     image_column_name: str | None = None,
     transform: Callable[[Any], Any] | None = None,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> Dataset:
     """Prepare the image dataset by converting images to RGB and applying transformations."""
     if (
@@ -315,7 +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,
+    num_proc: int | None = None,
 ) -> DataLoader[ImageInput]:
     """Creates a DataLoader with the image dataset prepared using the explicit transformation.
 
@@ -341,14 +341,14 @@ def _create_image_dataloader(
         batch_size=batch_size,
         collate_fn=collate_fn,
         shuffle=False,
-        num_workers=num_proc if num_proc > 1 else 0,
+        num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
     )
 
 
 def _create_text_queries_dataloader(
     dataset: Dataset,
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> DataLoader[QueryInput]:
     if not isinstance(dataset["text"][0], list):
         return _create_text_dataloader_for_queries(
@@ -368,7 +368,7 @@ def _create_queries_dataloader(
     task_metadata: TaskMetadata,
     input_column: str | None = None,
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> DataLoader[QueryInput | ImageInput]:
     """Create a dataloader for queries."""
     queries_type = task_metadata.get_modalities(PromptType.query)
@@ -393,7 +393,7 @@ def _create_document_dataloader(
     task_metadata: TaskMetadata,
     input_column: str | None = None,
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> DataLoader[CorpusInput | ImageInput]:
     """Create a dataloader for documents.
 
@@ -430,7 +430,7 @@ def create_dataloader(
     prompt_type: PromptType | None = None,
     input_column: str | None = None,
     batch_size: int = 32,
-    num_proc: int = 1,
+    num_proc: int | None = None,
     **kwargs: Any,
 ) -> DataLoader[BatchedInput]:
     """Create a dataloader from a dataset.
@@ -482,5 +482,5 @@ def create_dataloader(
     return DataLoader(
         dataset,
         batch_size=batch_size,
-        num_workers=num_proc if num_proc > 1 else 0,
+        num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
     )

mteb/_evaluators/any_sts_evaluator.py

@@ -66,7 +66,7 @@ class AnySTSEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> STSEvaluatorScores:
         logger.info("Running semantic similarity - Encoding samples (1/2)")
         embeddings1 = model.encode(

mteb/_evaluators/clustering_evaluator.py

@@ -45,7 +45,7 @@ class ClusteringEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> list[int]:
         data_loader = create_dataloader(
             self.dataset,

mteb/_evaluators/image/imagetext_pairclassification_evaluator.py

@@ -91,7 +91,7 @@ class ImageTextPairClassificationEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> list[torch.Tensor]:
         images = []
         if isinstance(self.images_column_names, str):
@@ -139,7 +139,7 @@ class ImageTextPairClassificationEvaluator(Evaluator):
             DataLoader(
                 CustomImageDataset(images),
                 collate_fn=_image_collate_fn,
-                num_workers=num_proc if num_proc > 1 else 0,
+                num_workers=num_proc if num_proc is not None and num_proc > 1 else 0,
             ),
             task_metadata=self.task_metadata,
             hf_subset=self.hf_subset,

mteb/_evaluators/pair_classification_evaluator.py

@@ -92,7 +92,7 @@ class PairClassificationEvaluator(Evaluator):
         self,
         model: EncoderProtocol,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> PairClassificationDistances:
         logger.info("Running pair classification - Encoding samples (1/2)")
         embeddings1 = model.encode(

mteb/_evaluators/retrieval_evaluator.py

@@ -55,7 +55,7 @@ class RetrievalEvaluator(Evaluator):
         self,
         search_model: SearchProtocol,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> RetrievalOutputType:
         logger.info("Running retrieval task - Indexing corpus...")
         search_model.index(

mteb/_evaluators/sklearn_evaluator.py

@@ -59,7 +59,9 @@ class SklearnEvaluator(Evaluator):
         self.evaluator_model = evaluator_model
 
     def create_dataloaders(
-        self, encode_kwargs: EncodeKwargs, num_proc: int
+        self,
+        encode_kwargs: EncodeKwargs,
+        num_proc: int | None,
     ) -> tuple[DataLoader[BatchedInput], DataLoader[BatchedInput]]:
         dataloader_train = create_dataloader(
             self.train_dataset,
@@ -83,7 +85,7 @@ class SklearnEvaluator(Evaluator):
         *,
         encode_kwargs: EncodeKwargs,
         test_cache: Array | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> tuple[NDArray[np.integer | np.floating], Array]:
         """Classification evaluation by training a sklearn classifier on the embeddings of the training set and evaluating on the embeddings of the test set.
 

mteb/_evaluators/text/bitext_mining_evaluator.py

@@ -41,7 +41,7 @@ class BitextMiningEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> dict[str, list[dict[str, float]]]:
         pair_elements = {p for pair in self.pairs for p in pair}
         if isinstance(self.sentences, Dataset):

mteb/_evaluators/text/summarization_evaluator.py

@@ -100,7 +100,7 @@ class SummarizationEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> 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]

mteb/_evaluators/zeroshot_classification_evaluator.py

@@ -48,7 +48,7 @@ class ZeroShotClassificationEvaluator(Evaluator):
         model: EncoderProtocol,
         *,
         encode_kwargs: EncodeKwargs,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> Array:
         dataloader = create_dataloader(
             self.dataset,

mteb/abstasks/abstask.py

@@ -116,7 +116,7 @@ class AbsTask(ABC):
             logger.warning(msg)
             warnings.warn(msg)
 
-    def dataset_transform(self, num_proc: int = 1, **kwargs: Any) -> None:
+    def dataset_transform(self, num_proc: int | None = None, **kwargs: Any) -> None:
         """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.
@@ -136,7 +136,7 @@ class AbsTask(ABC):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluates an MTEB compatible model on the task.
@@ -219,7 +219,7 @@ class AbsTask(ABC):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> ScoresDict:
         raise NotImplementedError(
@@ -324,7 +324,7 @@ class AbsTask(ABC):
             )  # only take the specified test split.
         return dataset_dict
 
-    def load_data(self, num_proc: int = 1, **kwargs: Any) -> None:
+    def load_data(self, num_proc: int | None = None, **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

mteb/abstasks/classification.py

@@ -138,7 +138,7 @@ class AbsTaskClassification(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
         """Evaluate a model on the classification task.
@@ -201,7 +201,7 @@ class AbsTaskClassification(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> FullClassificationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/clustering.py

@@ -169,7 +169,7 @@ class AbsTaskClustering(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> ScoresDict:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/clustering_legacy.py

@@ -95,7 +95,7 @@ class AbsTaskClusteringLegacy(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> ScoresDict:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/image/image_text_pair_classification.py

@@ -134,7 +134,7 @@ class AbsTaskImageTextPairClassification(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> ImageTextPairClassificationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/multilabel_classification.py

@@ -95,7 +95,7 @@ class AbsTaskMultilabelClassification(AbsTaskClassification):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> FullMultilabelClassificationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/pair_classification.py

@@ -97,7 +97,7 @@ class AbsTaskPairClassification(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> dict[str, float]:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/retrieval.py

@@ -148,7 +148,10 @@ class AbsTaskRetrieval(AbsTask):
             )
         )
 
-    def convert_v1_dataset_format_to_v2(self, num_proc: int) -> None:
+    def convert_v1_dataset_format_to_v2(
+        self,
+        num_proc: int | None,
+    ) -> None:
         """Convert dataset from v1 (from `self.queries`, `self.document`) format to v2 format (`self.dotaset`)."""
         # check if dataset is `v1` version
         if not hasattr(self, "queries"):
@@ -257,7 +260,7 @@ class AbsTaskRetrieval(AbsTask):
         if hasattr(self, "top_ranked"):
             del self.top_ranked
 
-    def load_data(self, num_proc: int = 1, **kwargs) -> None:
+    def load_data(self, num_proc: int | None = None, **kwargs) -> None:
         """Load the dataset for the retrieval task."""
         if self.data_loaded:
             return
@@ -301,7 +304,7 @@ class AbsTaskRetrieval(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> Mapping[HFSubset, ScoresDict]:
         """Evaluate the model on the retrieval task.
@@ -342,7 +345,7 @@ class AbsTaskRetrieval(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> ScoresDict:
         """Evaluate a model on a specific subset of the data.
@@ -473,7 +476,7 @@ class AbsTaskRetrieval(AbsTask):
         split: str,
         hf_subset: str | None = None,
         compute_overall: bool = False,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> RetrievalDescriptiveStatistics:
         self.convert_v1_dataset_format_to_v2(num_proc)
         if hf_subset and hf_subset in self.dataset:

mteb/abstasks/retrieval_dataset_loaders.py

@@ -78,7 +78,7 @@ class RetrievalDatasetLoader:
 
     def load(
         self,
-        num_proc: int = 1,
+        num_proc: int | None = None,
     ) -> RetrievalSplitData:
         """Loads the dataset split for the specified configuration.
 
@@ -128,7 +128,11 @@ class RetrievalDatasetLoader:
             f"Split {self.split} not found in {splits}. Please specify a valid split."
         )
 
-    def _load_dataset_split(self, config: str, num_proc: int) -> Dataset:
+    def _load_dataset_split(
+        self,
+        config: str,
+        num_proc: int | None,
+    ) -> Dataset:
         return load_dataset(
             self.hf_repo,
             config,
@@ -138,7 +142,10 @@ class RetrievalDatasetLoader:
             num_proc=num_proc,
         )
 
-    def _load_corpus(self, num_proc: int) -> CorpusDatasetType:
+    def _load_corpus(
+        self,
+        num_proc: int | None,
+    ) -> CorpusDatasetType:
         config = f"{self.config}-corpus" if self.config is not None else "corpus"
         logger.info("Loading corpus subset: %s", config)
 
@@ -151,7 +158,10 @@ class RetrievalDatasetLoader:
         logger.debug("Doc Example: %s", corpus_ds[0])
         return corpus_ds
 
-    def _load_queries(self, num_proc: int) -> QueryDatasetType:
+    def _load_queries(
+        self,
+        num_proc: int | None,
+    ) -> QueryDatasetType:
         config = f"{self.config}-queries" if self.config is not None else "queries"
         logger.info("Loading queries subset: %s", config)
 
@@ -168,7 +178,10 @@ class RetrievalDatasetLoader:
 
         return queries_ds
 
-    def _load_qrels(self, num_proc: int) -> RelevantDocumentsType:
+    def _load_qrels(
+        self,
+        num_proc: int | None,
+    ) -> RelevantDocumentsType:
         config = f"{self.config}-qrels" if self.config is not None else "default"
 
         logger.info("Loading qrels subset: %s", config)
@@ -203,7 +216,10 @@ class RetrievalDatasetLoader:
         logger.info("Loaded %d %s qrels.", len(qrels_dict), self.split.upper())
         return qrels_dict
 
-    def _load_top_ranked(self, num_proc: int) -> TopRankedDocumentsType:
+    def _load_top_ranked(
+        self,
+        num_proc: int | None,
+    ) -> TopRankedDocumentsType:
         config = (
             f"{self.config}-top_ranked" if self.config is not None else "top_ranked"
         )
@@ -226,7 +242,10 @@ class RetrievalDatasetLoader:
         logger.info(f"Top ranked loaded: {len(top_ranked_ds)}")
         return top_ranked_dict
 
-    def _load_instructions(self, num_proc: int) -> InstructionDatasetType:
+    def _load_instructions(
+        self,
+        num_proc: int | None,
+    ) -> InstructionDatasetType:
         config = (
             f"{self.config}-instruction" if self.config is not None else "instruction"
         )
@@ -246,7 +265,7 @@ class RetrievalDatasetLoader:
 def _combine_queries_with_instructions_datasets(
     queries_dataset: QueryDatasetType,
     instruction_dataset: InstructionDatasetType | dict[str, str],
-    num_proc: int,
+    num_proc: int | None,
 ) -> Dataset:
     if isinstance(instruction_dataset, Dataset):
         instruction_to_query_idx = {

mteb/abstasks/sts.py

@@ -118,7 +118,7 @@ class AbsTaskSTS(AbsTask):
         hf_split: str,
         hf_subset: str,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> STSMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/text/bitext_mining.py

@@ -82,7 +82,7 @@ class AbsTaskBitextMining(AbsTask):
         *,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs: Any,
     ) -> dict[HFSubset, ScoresDict]:
         """Added load for "parallel" datasets"""
@@ -155,7 +155,7 @@ class AbsTaskBitextMining(AbsTask):
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
         parallel: bool = False,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> BitextMiningMetrics | dict[str, BitextMiningMetrics]:
         pairs = self._get_pairs(parallel)

mteb/abstasks/text/reranking.py

@@ -34,7 +34,7 @@ class AbsTaskReranking(AbsTaskRetrieval):
         For dataformat and other information, see [AbsTaskRetrieval][mteb.abstasks.retrieval.AbsTaskRetrieval].
     """
 
-    def load_data(self, num_proc: int = 1, **kwargs) -> None:
+    def load_data(self, num_proc: int | None = None, **kwargs) -> None:
         """Load the dataset."""
         if self.data_loaded:
             return

mteb/abstasks/text/summarization.py

@@ -94,7 +94,7 @@ class AbsTaskSummarization(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> SummarizationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/abstasks/zeroshot_classification.py

@@ -127,7 +127,7 @@ class AbsTaskZeroShotClassification(AbsTask):
         hf_subset: str,
         encode_kwargs: EncodeKwargs,
         prediction_folder: Path | None = None,
-        num_proc: int = 1,
+        num_proc: int | None = None,
         **kwargs,
     ) -> ZeroShotClassificationMetrics:
         if not isinstance(model, EncoderProtocol):

mteb/evaluate.py

@@ -91,7 +91,7 @@ def _evaluate_task(
     encode_kwargs: EncodeKwargs,
     prediction_folder: Path | None,
     public_only: bool | None,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> TaskResult | TaskError:
     """The core logic to run a model on a given task. See `evaluate` for more details.
 
@@ -282,7 +282,7 @@ def evaluate(
     prediction_folder: Path | str | None = None,
     show_progress_bar: bool = True,
     public_only: bool | None = None,
-    num_proc: int = 1,
+    num_proc: int | None = None,
 ) -> ModelResult:
     """This function runs a model on a given task and returns the results.
 

mteb/models/model_implementations/bm25.py

@@ -54,7 +54,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
             hf_split: str,
             hf_subset: str,
             encode_kwargs: EncodeKwargs,
-            num_proc: int = 1,
+            num_proc: int | None = None,
         ) -> None:
             logger.info("Encoding Corpus...")
             corpus_texts = [
@@ -81,7 +81,7 @@ def bm25_loader(model_name, **kwargs) -> SearchProtocol:
             top_k: int,
             encode_kwargs: EncodeKwargs,
             top_ranked: TopRankedDocumentsType | None = None,
-            num_proc: int = 1,
+            num_proc: int | None = None,
         ) -> RetrievalOutputType:
             logger.info("Encoding Queries...")
             query_ids = list(queries["id"])