Fix selection of DRES/DRPES

mteb/abstasks/AbsTaskRetrieval.py

@@ -39,6 +39,7 @@ def evaluate(
         batch_size=128,
         corpus_chunk_size=None,
         score_function="cos_sim",
+        parallel_retrieval=False,
         **kwargs
     ):
         try:
@@ -47,12 +48,12 @@ def evaluate(
             raise Exception("Retrieval tasks require beir package. Please install it with `pip install mteb[beir]`")
 
         if not self.data_loaded:
-            self.load_data()
+            self.load_data(parallel_retrieval=parallel_retrieval)
 
         corpus, queries, relevant_docs = self.corpus[split], self.queries[split], self.relevant_docs[split]
         model = model if self.is_dres_compatible(model) else DRESModel(model)
 
-        if os.getenv("RANK", None) is None:
+        if not parallel_retrieval:
             # Non-distributed
             from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
             model = DRES(

mteb/abstasks/BeIRTask.py

@@ -11,7 +11,7 @@ class BeIRTask(AbsTask):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
-    def load_data(self, eval_splits=None, **kwargs):
+    def load_data(self, eval_splits=None, parallel_retrieval=False, **kwargs):
         """
         Load dataset from BeIR benchmark. TODO: replace with HF hub once datasets are moved there
         """
@@ -23,7 +23,7 @@ def load_data(self, eval_splits=None, **kwargs):
         USE_HF_DATASETS = False
 
         # TODO @nouamane: move non-distributed to `HFDataLoader`
-        if os.getenv("RANK", None) is not None:
+        if parallel_retrieval:
             if self.description["beir_name"].startswith("cqadupstack"):
                 raise ImportError("CQADupstack is incompatible with BEIR's HFDataLoader in a distributed setting")
             from beir.datasets.data_loader_hf import HFDataLoader 

mteb/evaluation/MTEB.py

@@ -258,7 +258,7 @@ def run(
 
                 # load data
                 logger.info(f"Loading dataset for {task.description['name']}")
-                task.load_data(eval_splits=task_eval_splits)
+                task.load_data(eval_splits=task_eval_splits, **kwargs)
 
                 # run evaluation
                 task_results = {