Merge pull request #7 from jina-ai/fix-single-gpu

Allow running on single gpu

mteb/abstasks/AbsTaskRetrieval.py

@@ -52,7 +52,7 @@ def evaluate(
         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 os.getenv("RANK", None) is None or os.getenv("MTEB_SINGLE_GPU", "false").lower() == "true":
             # Non-distributed
             from beir.retrieval.search.dense import DenseRetrievalExactSearch as DRES
             model = DRES(

mteb/abstasks/BeIRTask.py

@@ -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 os.getenv("RANK", None) is not None and os.getenv("MTEB_SINGLE_GPU", "false").lower() == "false":
             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