[utils] - feat: experiment tracker

AGENTS.md

@@ -26,6 +26,9 @@
 - Naming: modules/files `snake_case.py`; classes `CapWords`; functions/vars `snake_case`.
 - Prefer type hints and docstrings; keep public APIs stable under `bert_squeeze/`.
 
+### Typing
+- Do not use type `Any`; be as strict as possible on type
+
 ## Testing Guidelines
 - Use `pytest`; place new tests under `tests/` as `test_*.py`.
 - Keep tests deterministic (avoid network); use `tests/fixtures/` for small assets.

README.md

@@ -46,6 +46,9 @@ uv sync
 
 You are all set!
 
+Optional: install Aim support with `uv sync --extra aim` and configure
+`logger_kwargs={"_target_": "lightning.pytorch.loggers.AimLogger"}` in an assistant.
+
 # Quickstarts
 
 You can find a bunch of examples on how to use the library to simply train models or perform optimization techniques 

bert_squeeze/assistants/distil_assistant.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 import lightning.pytorch as pl
 import torch.nn
@@ -38,17 +38,17 @@ class DistilAssistant(object):
     Args:
         name (str):
             name of the base model to fine-tune
-        general_kwargs (Dict[str, Any]):
+        general_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'general' configuration
-        train_kwargs (Dict[str, Any]):
+        train_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'train' configuration
-        student_kwargs (Dict[str, Any]):
+        student_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'model.student' configuration
-        teacher_kwargs (Dict[str, Any]):
+        teacher_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'model.teacher' configuration
-        data_kwargs (Dict[str, Any]):
+        data_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'data' configuration
-        logger_kwargs (Dict[str, Any]):
+        logger_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'logger' configuration
         callbacks (List[Callback]):
             list of callbacks to use during training
@@ -68,13 +68,13 @@ class DistilAssistant(object):
     def __init__(
         self,
         name: str,
-        general_kwargs: Dict[str, Any] = None,
-        train_kwargs: Dict[str, Any] = None,
-        student_kwargs: Dict[str, Any] = None,
-        teacher_kwargs: Dict[str, Any] = {},
-        data_kwargs: Dict[str, Any] = None,
-        logger_kwargs: Dict[str, Any] = None,
-        callbacks: List[Callback] = None,
+        general_kwargs: Optional[Dict[str, object]] = None,
+        train_kwargs: Optional[Dict[str, object]] = None,
+        student_kwargs: Optional[Dict[str, object]] = None,
+        teacher_kwargs: Dict[str, object] = {},
+        data_kwargs: Optional[Dict[str, object]] = None,
+        logger_kwargs: Optional[Dict[str, object]] = None,
+        callbacks: Optional[List[Callback]] = None,
     ):
         conf = OmegaConf.load(
             resource_filename(
@@ -93,6 +93,11 @@ def __init__(
             [general_kwargs, train_kwargs, data_kwargs, logger_kwargs, callbacks],
         ):
             if kws is not None:
+                base = conf.get(name)
+                if base is None:
+                    conf[name] = kws
+                    continue
+
                 if "_target_" in kws and kws["_target_"] != conf[name]["_target_"]:
                     del conf[name]
                     conf[name] = kws
@@ -137,7 +142,7 @@ def student_config(self) -> DictConfig:
         return self._model_conf["student"]
 
     @property
-    def model(self) -> Any:
+    def model(self) -> pl.LightningModule:
         """"""
         if self._model is None:
             self.model = instantiate(self._model_conf)
@@ -159,27 +164,27 @@ def model(self) -> Any:
         return self._model
 
     @model.setter
-    def model(self, value: Any) -> None:
+    def model(self, value: pl.LightningModule) -> None:
         self._model = value
 
     @property
-    def student(self) -> Any:
+    def student(self) -> Optional[Union[pl.LightningModule, torch.nn.Module]]:
         """"""
         if self._model is None:
             logging.warning("The Distiller has not been instantiated.")
             return None
         return self.model.student
 
     @property
-    def teacher(self) -> Any:
+    def teacher(self) -> Optional[Union[pl.LightningModule, torch.nn.Module]]:
         """"""
         if self._model is None:
             logging.warning("The Distiller has not been instantiated.")
             return None
         return self.model.teacher
 
     @property
-    def data(self) -> Any:
+    def data(self) -> pl.LightningDataModule:
         """"""
         if self._data is None:
             data = instantiate(self._data_conf, _recursive_=True)
@@ -189,7 +194,7 @@ def data(self) -> Any:
         return self._data
 
     @data.setter
-    def data(self, value: Any) -> None:
+    def data(self, value: pl.LightningDataModule) -> None:
         """"""
         self._data = value
 
@@ -221,7 +226,7 @@ def callbacks(self) -> Optional[List[Callback]]:
         return self._callbacks
 
     @callbacks.setter
-    def callbacks(self, value) -> None:
+    def callbacks(self, value: Optional[List[Callback]]) -> None:
         """"""
         self._callbacks = value
 

bert_squeeze/assistants/train_assistant.py

@@ -1,13 +1,15 @@
 import logging
 import os
-from typing import Any, Dict, List
+from typing import Dict, List, Optional
 
+import lightning.pytorch as pl
 from hydra.utils import instantiate
 from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.loggers import Logger, TensorBoardLogger
 from omegaconf import OmegaConf
 from pkg_resources import resource_filename
-from pydantic.utils import deep_update
+
+from bert_squeeze.utils.utils_fct import deep_update
 
 CONFIG_MAPPER = {
     "lr": "train_lr.yaml",
@@ -38,15 +40,15 @@ class TrainAssistant(object):
     Args:
         name (str):
             name of the base model to fine-tune
-        general_kwargs (Dict[str, Any]):
+        general_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'general' configuration
-        train_kwargs (Dict[str, Any]):
+        train_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'train' configuration
-        model_kwargs (Dict[str, Any]):
+        model_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'model' configuration
-        data_kwargs (Dict[str, Any]):
+        data_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'data' configuration
-        logger_kwargs (Dict[str, Any]):
+        logger_kwargs (Dict[str, object]):
             keyword arguments that can be added or overwrite the default 'logger' configuration
         callbacks (List[Callback]):
             list of callbacks to use during training
@@ -55,12 +57,12 @@ class TrainAssistant(object):
     def __init__(
         self,
         name: str,
-        general_kwargs: Dict[str, Any] = None,
-        train_kwargs: Dict[str, Any] = None,
-        model_kwargs: Dict[str, Any] = None,
-        data_kwargs: Dict[str, Any] = None,
-        logger_kwargs: Dict[str, Any] = None,
-        callbacks: List[Callback] = None,
+        general_kwargs: Optional[Dict[str, object]] = None,
+        train_kwargs: Optional[Dict[str, object]] = None,
+        model_kwargs: Optional[Dict[str, object]] = None,
+        data_kwargs: Optional[Dict[str, object]] = None,
+        logger_kwargs: Optional[Dict[str, object]] = None,
+        callbacks: Optional[List[Callback]] = None,
     ):
         try:
             conf = OmegaConf.load(
@@ -100,7 +102,8 @@ def __init__(
             ],
         ):
             if kws is not None:
-                conf[name] = deep_update(conf[name], kws)
+                base = conf.get(name)
+                conf[name] = kws if base is None else deep_update(base, kws)
 
         self.name = name
         self.general = conf["general"]
@@ -110,25 +113,25 @@ def __init__(
         self._logger_conf = conf.get("logger")
         self._callbacks_conf = conf.get("callbacks", [])
 
-        self._model = None
-        self._data = None
+        self._model: Optional[pl.LightningModule] = None
+        self._data: Optional[pl.LightningDataModule] = None
         self._logger = None
         self._callbacks = None
 
     @property
-    def model(self) -> Any:
+    def model(self) -> pl.LightningModule:
         """"""
         if self._model is None:
             self.model = instantiate(self._model_conf)
         return self._model
 
     @model.setter
-    def model(self, value: Any) -> None:
+    def model(self, value: pl.LightningModule) -> None:
         """"""
         self._model = value
 
     @property
-    def data(self) -> Any:
+    def data(self) -> pl.LightningDataModule:
         """"""
         if self._data is None:
             data = instantiate(self._data_conf)
@@ -138,7 +141,7 @@ def data(self) -> Any:
         return self._data
 
     @data.setter
-    def data(self, value: Any) -> None:
+    def data(self, value: pl.LightningDataModule) -> None:
         """"""
         self._data = value
 

bert_squeeze/distillation/base_distiller.py

@@ -7,6 +7,7 @@
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 from transformers import AdamW
 
+from ..utils.experiment_logging import ExperimentLogger
 from ..utils.optimizers import BertAdam
 from ..utils.types import DistillationLoss
 
@@ -229,7 +230,7 @@ def log_eval_report(self, *args) -> None:
         """
         results = self.s_valid_scorer.to_dict()
         table = self.s_valid_scorer.get_table(results)
-        self.logger.experiment.add_text("eval/report", table)
+        ExperimentLogger.from_module(self).add_text("eval/report", table)
 
         # logging losses to neptune
         logging_loss = {
@@ -240,5 +241,5 @@ def log_eval_report(self, *args) -> None:
 
         # logging other metrics
         for key, value in results.items():
-            if not isinstance(value, list) and not isinstance(value, np.ndarray):
+            if not isinstance(value, (list, np.ndarray)):
                 self.log_dict({f"eval/{key}": value})

bert_squeeze/distillation/sequence_classification_distiller.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, List, Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 import lightning.pytorch as pl
 import matplotlib.pyplot as plt
@@ -13,6 +13,7 @@
 from transformers.modeling_outputs import SequenceClassifierOutput
 
 from bert_squeeze.distillation.base_distiller import BaseDistiller
+from bert_squeeze.utils.experiment_logging import ExperimentLogger
 from bert_squeeze.utils.losses import LabelSmoothingLoss
 from bert_squeeze.utils.losses.distillation_losses import KLDivLoss
 from bert_squeeze.utils.scorers import BaseSequenceClassificationScorer
@@ -71,7 +72,7 @@ def _set_objectives(self) -> None:
                 "You are using label smoothing and the smoothing parameteris set to 0.0."
             )
         elif objective == "weighted" and all(
-            [w == 1.0 for w in self.params.get("class_weights", None)]
+            w == 1.0 for w in self.params.get("class_weights", [])
         ):
             logging.warning(
                 "You are using a weighted CrossEntropy but the class"
@@ -103,13 +104,15 @@ def _set_scorers(self) -> None:
         self.s_valid_scorer = BaseSequenceClassificationScorer(self.labels)
         self.s_test_scorer = BaseSequenceClassificationScorer(self.labels)
 
-    def get_teacher_logits(self, batch: Dict[str, torch.Tensor]) -> Any:
+    def get_teacher_logits(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
         raise NotImplementedError()
 
-    def get_student_logits(self, batch: Dict[str, torch.Tensor]) -> Any:
+    def get_student_logits(
+        self, batch: Dict[str, torch.Tensor]
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         raise NotImplementedError()
 
-    def log_eval_report(self, probs: List[np.array]) -> None:
+    def log_eval_report(self, probs: List[np.ndarray]) -> None:
         """
         Method that logs an evaluation report.
 
@@ -123,11 +126,12 @@ def log_eval_report(self, probs: List[np.array]) -> None:
         super().log_eval_report()
 
         # logging probability distributions
-        for i in range(len(probs)):
+        exp_logger = ExperimentLogger.from_module(self)
+        for i, prob in enumerate(probs):
             fig = plt.figure(figsize=(15, 15))
-            sns.distplot(probs[i], kde=False, bins=100)
-            plt.title("Probability boxplot for label {}".format(i))
-            self.logger.experiment.add_figure("eval/dist_label_{}".format(i), fig)
+            sns.distplot(prob, kde=False, bins=100)
+            plt.title(f"Probability boxplot for label {i}")
+            exp_logger.add_figure(f"eval/dist_label_{i}", fig)
             plt.close("all")
 
 

bert_squeeze/models/base_lt_module.py

@@ -19,6 +19,7 @@
     AutoModelForSequenceClassification,
 )
 
+from ..utils.experiment_logging import ExperimentLogger
 from ..utils.losses import LabelSmoothingLoss
 from ..utils.optimizers import BertAdam
 from ..utils.scorers import BaseSequenceClassificationScorer, LMScorer, Scorer
@@ -307,7 +308,8 @@ def log_eval_report(self, *args) -> None:
             table = self.valid_scorer.get_table(eval_report)
         except TypeError:
             table = self.valid_scorer.get_table()
-        self.logger.experiment.add_text("eval/report", table)
+        exp_logger = ExperimentLogger.from_module(self)
+        exp_logger.add_text("eval/report", table)
 
         logging_loss = {}
         for key, values in self.valid_scorer.losses.items():
@@ -325,8 +327,8 @@ def log_eval_report(self, *args) -> None:
         for metric, value in eval_report.items():
             if isinstance(value, dict):
                 self.log_dict({f"eval/{metric}/{key}": v for key, v in value.items()})
-            elif not isinstance(value, list) and not isinstance(value, np.ndarray):
-                self.log("eval/{}".format(metric), value)
+            elif not isinstance(value, (list, np.ndarray)):
+                self.log(f"eval/{metric}", value)
 
 
 class BaseSequenceClassificationTransformerModule(BaseTransformerModule):
@@ -408,7 +410,7 @@ def _set_objective(self) -> None:
             logging.warning(
                 "You are using label smoothing and the smoothing parameteris set to 0.0."
             )
-        elif objective == "weighted" and all([w == 1.0 for w in self.class_weights]):
+        elif objective == "weighted" and all(w == 1.0 for w in self.class_weights):
             logging.warning(
                 "You are using a weighted CrossEntropy but the class"
                 "weights are all equal to 1.0."
@@ -461,11 +463,12 @@ def log_eval_report(self, probs: np.array) -> None:
         """
         super().log_eval_report()
 
+        exp_logger = ExperimentLogger.from_module(self)
         for i in range(probs.shape[1]):
             fig = plt.figure(figsize=(15, 15))
             sns.histplot(probs[:, i], bins=100)
-            plt.title("Probability boxplot for label {}".format(i))
-            self.logger.experiment.add_figure("eval/dist_label_{}".format(i), fig)
+            plt.title(f"Probability boxplot for label {i}")
+            exp_logger.add_figure(f"eval/dist_label_{i}", fig)
             plt.close("all")