feat: add DISABLE_EMISSIONS_TRACKING (#191)

* feat: add DISABLE_EMISSIONS_TRACKING

* try to fix docs error

* Update optimizer_config.schema.json

* another attempt

* Update optimizer_config.schema.json

* i give up for now

* Update optimizer_config.schema.json

---------

Co-authored-by: voorhs <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

autointent/configs/_transformers.py

@@ -115,7 +115,10 @@ def get_prompt_type(self, prompt_type: TaskTypeEnum | None) -> str | None:  # no
 
 
 class CrossEncoderConfig(HFModelConfig):
-    model_name: str = Field("cross-encoder/ms-marco-MiniLM-L-6-v2", description="Name of the hugging face model.")
+    model_name: str = Field("cross-encoder/ms-marco-MiniLM-L6-v2", description="Name of the hugging face model.")
     train_head: bool = Field(
         False, description="Whether to train the head of the model. If False, LogReg will be trained."
     )
+    tokenizer_config: TokenizerConfig = Field(
+        default_factory=lambda: TokenizerConfig(max_length=512)
+    )  # this is because sentence-transformers doesn't allow you to customize tokenizer settings properly

autointent/modules/scoring/_dnnc/dnnc.py

@@ -36,7 +36,7 @@ class DNNCScorer(BaseScorer):
         utterances = ["what is your name?", "how are you?"]
         labels = [0, 1]
         scorer = DNNCScorer(
-            cross_encoder_config="cross-encoder/ms-marco-MiniLM-L-6-v2",
+            cross_encoder_config="cross-encoder/ms-marco-MiniLM-L6-v2",
             embedder_config="sergeyzh/rubert-tiny-turbo",
             k=5,
         )

autointent/nodes/emissions_tracker.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import os
 
 from codecarbon import EmissionsTracker as CodeCarbonTracker  # type: ignore[import-untyped]
 from codecarbon.output import EmissionsData  # type: ignore[import-untyped]
@@ -20,22 +21,31 @@ def __init__(self, project_name: str, measure_power_secs: int = 1) -> None:
             measure_power_secs: How often to measure power consumption in seconds.
         """
         self._logger = logger
-        self.tracker = CodeCarbonTracker(project_name=project_name, measure_power_secs=measure_power_secs)
+        self._enabled = int(os.getenv("TRACK_EMISSIONS", "0"))
+        if self._enabled:
+            self.tracker = CodeCarbonTracker(project_name=project_name, measure_power_secs=measure_power_secs)
+        else:
+            self._logger.info("Emissions tracking is enabled via TRACK_EMISSIONS environment variable")
+            self.tracker = None
 
     def start_task(self, task_name: str) -> None:
         """Start tracking emissions for a specific task.
 
         Args:
             task_name: Name of the task to track emissions for.
         """
-        self.tracker.start_task(task_name)
+        if self._enabled:
+            self.tracker.start_task(task_name)
 
     def stop_task(self) -> dict[str, float]:
         """Stop tracking emissions and return the emissions data.
 
         Returns:
             Dictionary containing emissions metrics.
         """
+        if not self._enabled:
+            return {}
+
         emissions_data = self.tracker.stop_task()
         _ = self.tracker.stop()
         return self._process_metrics(emissions_data)

docs/optimizer_config.schema.json

@@ -4,7 +4,7 @@
             "additionalProperties": false,
             "properties": {
                 "model_name": {
-                    "default": "cross-encoder/ms-marco-MiniLM-L-6-v2",
+                    "default": "cross-encoder/ms-marco-MiniLM-L6-v2",
                     "description": "Name of the hugging face model.",
                     "title": "Model Name",
                     "type": "string"
@@ -424,11 +424,11 @@
         "cross_encoder_config": {
             "$ref": "#/$defs/CrossEncoderConfig",
             "default": {
-                "model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2",
+                "model_name": "cross-encoder/ms-marco-MiniLM-L6-v2",
                 "batch_size": 32,
                 "device": null,
                 "tokenizer_config": {
-                    "max_length": null,
+                    "max_length": 512,
                     "padding": true,
                     "truncation": true
                 },

tests/_transformers/test_nli_transformer.py

@@ -13,7 +13,7 @@ def data_handler():
 
 
 def test_nli_transformer_predict_without_trained_head(data_handler):
-    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2", "train_head": True})
+    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L6-v2", "train_head": True})
     with pytest.raises(ValueError, match="Classifier is not trained yet"):
         model.predict(data_handler.train_utterances(0))
 
@@ -48,7 +48,7 @@ def check_ranking(ranked, labels):
 
 
 def test_nli_transformer_predict_with_train_head(data_handler):
-    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2", "train_head": True})
+    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L6-v2", "train_head": True})
     texts = data_handler.train_utterances(0)
     labels = data_handler.train_labels(0)
     model.fit(texts, labels)
@@ -60,7 +60,7 @@ def test_nli_transformer_predict_with_train_head(data_handler):
 
 
 def test_nli_transformer_predict_default(data_handler):
-    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2", "train_head": False})
+    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L6-v2", "train_head": False})
     texts = data_handler.train_utterances(0)
     labels = data_handler.train_labels(0)
     predicted = model.predict(build_pairs(texts))
@@ -71,7 +71,7 @@ def test_nli_transformer_predict_default(data_handler):
 
 
 def test_nli_transformer_predict_default_with_fit(data_handler):
-    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2", "train_head": False})
+    model = Ranker(cross_encoder_config={"model_name": "cross-encoder/ms-marco-MiniLM-L6-v2", "train_head": False})
     texts = data_handler.train_utterances(0)
     labels = data_handler.train_labels(0)
     model.fit(texts, labels)

tests/assets/configs/description.yaml

@@ -14,7 +14,7 @@
       embedder_config:
         - model_name: sentence-transformers/all-MiniLM-L6-v2
       cross_encoder_config:
-        - model_name: cross-encoder/ms-marco-MiniLM-L-6-v2
+        - model_name: cross-encoder/ms-marco-MiniLM-L6-v2
       encoder_type: [cross, bi]
 - node_type: decision
   target_metric: decision_accuracy

tests/assets/configs/full_training.yaml

@@ -30,4 +30,4 @@ embedder_config:
   use_cache: true
 cross_encoder_config:
   batch_size: 32
-  model_name: cross-encoder/ms-marco-MiniLM-L-6-v2
+  model_name: cross-encoder/ms-marco-MiniLM-L6-v2

tests/assets/configs/multiclass.yaml

@@ -15,7 +15,7 @@
     - module_name: linear
     - module_name: dnnc
       cross_encoder_config:
-        - model_name: cross-encoder/ms-marco-MiniLM-L-6-v2
+        - model_name: cross-encoder/ms-marco-MiniLM-L6-v2
           train_head: true
         - avsolatorio/GIST-small-Embedding-v0
       k: [1, 3]
@@ -25,7 +25,7 @@
       m: [ 2, 3 ]
       use_crosencoder_scores: [true, false]
       cross_encoder_config:
-        - cross-encoder/ms-marco-MiniLM-L-6-v2
+        - cross-encoder/ms-marco-MiniLM-L6-v2
     - module_name: sklearn
       clf_name: [RandomForestClassifier]
       n_estimators: [5, 10]

tests/assets/configs/multilabel.yaml

@@ -21,7 +21,7 @@
       use_crosencoder_scores: [true, false]
       m: [ 2, 3 ]
       cross_encoder_config:
-        - model_name: cross-encoder/ms-marco-MiniLM-L-6-v2
+        - model_name: cross-encoder/ms-marco-MiniLM-L6-v2
     - module_name: sklearn
       clf_name: [RandomForestClassifier]
       n_estimators: [5, 10]

tests/configs/test_combined_config.py

@@ -17,7 +17,7 @@ def valid_optimizer_config():
                 {
                     "module_name": "dnnc",
                     "cross_encoder_config": [
-                        {"model_name": "cross-encoder/ms-marco-MiniLM-L-6-v2", "train_head": True},
+                        {"model_name": "cross-encoder/ms-marco-MiniLM-L6-v2", "train_head": True},
                         {"model_name": "avsolatorio/GIST-small-Embedding-v0", "train_head": False},
                     ],
                     "k": [1, 3],
@@ -63,7 +63,7 @@ def test_invalid_optimizer_config_missing_field():
             "node_type": "scoring",
             # Missing "target_metric"
             "search_space": [
-                {"module_name": "dnnc", "cross_encoder_name": ["cross-encoder/ms-marco-MiniLM-L-6-v2"], "k": [1, 3]}
+                {"module_name": "dnnc", "cross_encoder_name": ["cross-encoder/ms-marco-MiniLM-L6-v2"], "k": [1, 3]}
             ],
         }
     ]
@@ -80,7 +80,7 @@ def test_invalid_optimizer_config_wrong_type():
         "search_space": [
             {
                 "module_name": "dnnc",
-                "cross_encoder_name": "cross-encoder/ms-marco-MiniLM-L-6-v2",  # Should be a list
+                "cross_encoder_name": "cross-encoder/ms-marco-MiniLM-L6-v2",  # Should be a list
                 "k": "wrong_type",  # Should be a list of integers
                 "train_head": "true",  # Should be a boolean, not a string
             }