feat: enable codecarbon by default (#1428)

* enable codecarbon by default * lint * update flag * add allow_multiple_runs param * make lint * add warning * lint * negate the flag --------- Co-authored-by: Isaac Chung <[email protected]>
embeddings-benchmark · Nov 15, 2024 · 70a3ff2 · 70a3ff2
1 parent 5470c88
commit 70a3ff2
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 15 deletions.
diff --git a/mteb/cli.py b/mteb/cli.py
@@ -142,12 +142,14 @@ def run(args: argparse.Namespace) -> None:
         args.save_predictions if hasattr(args, "save_predictions") else False
     )
 
+    enable_co2_tracker = not args.disable_co2_tracker
+
     eval.run(
         model,
         verbosity=args.verbosity,
         output_folder=args.output_folder,
         eval_splits=args.eval_splits,
-        co2_tracker=args.co2_tracker,
+        co2_tracker=enable_co2_tracker,
         overwrite_results=args.overwrite,
         encode_kwargs=encode_kwargs,
         save_predictions=save_predictions,
@@ -263,10 +265,10 @@ def add_run_parser(subparsers) -> None:
         "-v", "--verbosity", type=int, default=2, help="Verbosity level"
     )
     parser.add_argument(
-        "--co2_tracker",
-        type=bool,
+        "--disable_co2_tracker",
+        action="store_true",
         default=False,
-        help="Enable CO₂ tracker, disabled by default",
+        help="Disable CO₂ tracker, enabled by default",
     )
     parser.add_argument(
         "--eval_splits",

diff --git a/mteb/evaluation/MTEB.py b/mteb/evaluation/MTEB.py
@@ -13,6 +13,7 @@
 from typing import Any
 
 import datasets
+from codecarbon import EmissionsTracker
 from sentence_transformers import SentenceTransformer
 
 from mteb.encoder_interface import Encoder
@@ -322,7 +323,7 @@ def run(
         eval_splits=None,
         overwrite_results: bool = False,
         raise_error: bool = True,
-        co2_tracker: bool = False,
+        co2_tracker: bool = True,
         encode_kwargs: dict[str, Any] = {},
         **kwargs,
     ) -> list[TaskResult]:
@@ -418,15 +419,14 @@ def run(
                 kg_co2_emissions: int | None = 0 if co2_tracker else None
                 for split in task_eval_splits:
                     if co2_tracker:
-                        try:
-                            from codecarbon import EmissionsTracker
-                        except ImportError:
-                            raise ImportError(
-                                "To use the CO2 emissions tracker, please install codecarbon using 'pip install codecarbon'"
-                            )
-
+                        logger.warning(
+                            "Evaluating multiple MTEB runs simultaniously will produce incorrect CO₂ results"
+                        )
                         with EmissionsTracker(
-                            save_to_file=False, save_to_api=False, logging_logger=logger
+                            save_to_file=False,
+                            save_to_api=False,
+                            logging_logger=logger,
+                            allow_multiple_runs=True,
                         ) as tracker:
                             results, tick, tock = self._run_eval(
                                 task,

diff --git a/pyproject.toml b/pyproject.toml
@@ -40,6 +40,7 @@ dependencies = [
     "typing_extensions>=0.0.0",
     "eval_type_backport>=0.0.0",
     "polars>=0.20.22",
+    "codecarbon>=2.0.0",
 ]
 
 
@@ -54,7 +55,6 @@ mteb = "mteb.cli:main"
 [project.optional-dependencies]
 dev = ["ruff==0.6.4", # locked so we don't get PRs which fail only due to a lint update
 "pytest", "pytest-xdist", "pytest-coverage"]
-codecarbon = ["codecarbon"]
 speedtask = ["GPUtil>=1.4.0", "psutil>=5.9.8"]
 peft = ["peft>=0.11.0"]
 leaderboard = ["gradio>=5.5.0", "gradio_rangeslider>=0.0.8"]

diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -62,7 +62,7 @@ def test_run_task(
         task_types=None,
         languages=None,
         batch_size=None,
-        co2_tracker=None,
+        disable_co2_tracker=None,
         overwrite=True,
         eval_splits=None,
         benchmarks=None,