Merge branch 'dev' into out_dim_dynamic

Author: sfluegel05

.github/workflows/test.yml

@@ -9,19 +9,30 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
 
     steps:
       - uses: actions/checkout@v4
+
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           python -m pip install --upgrade pip setuptools wheel
           python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
           python -m pip install -e .
-      - name: Display Python version
-        run: python -m unittest discover -s tests/unit
+
+      - name: Display Python & Installed Packages
+        run: |
+          python --version
+          pip freeze
+
+      - name: Run Unit Tests
+        run: python -m unittest discover -s tests/unit -v
+        env:
+          ACTIONS_STEP_DEBUG: true  # Enable debug logs
+          ACTIONS_RUNNER_DEBUG: true # Additional debug logs from Github Actions itself

chebai/cli.py

@@ -58,9 +58,6 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser):
         parser.link_arguments(
             "data.num_of_labels", "trainer.callbacks.init_args.num_labels"
         )
-        parser.link_arguments(
-            "data", "model.init_args.criterion.init_args.data_extractor"
-        )
 
     @staticmethod
     def subcommands() -> Dict[str, Set[str]]:

chebai/loss/bce_weighted.py

@@ -5,14 +5,16 @@
 import torch
 
 from chebai.preprocessing.datasets.base import XYBaseDataModule
+from chebai.preprocessing.datasets.chebi import _ChEBIDataExtractor
 from chebai.preprocessing.datasets.pubchem import LabeledUnlabeledMixed
 
 
 class BCEWeighted(torch.nn.BCEWithLogitsLoss):
     """
     BCEWithLogitsLoss with weights automatically computed according to the beta parameter.
+    If beta is None or data_extractor is None, the loss is unweighted.
 
-    This class computes weights based on the formula from the paper:
+    This class computes weights based on the formula from the paper by Cui et al. (2019):
     https://openaccess.thecvf.com/content_CVPR_2019/papers/Cui_Class-Balanced_Loss_Based_on_Effective_Number_of_Samples_CVPR_2019_paper.pdf
 
     Args:
@@ -24,13 +26,17 @@ def __init__(
         self,
         beta: Optional[float] = None,
         data_extractor: Optional[XYBaseDataModule] = None,
+        **kwargs,
     ):
         self.beta = beta
         if isinstance(data_extractor, LabeledUnlabeledMixed):
             data_extractor = data_extractor.labeled
         self.data_extractor = data_extractor
-
-        super().__init__()
+        assert (
+            isinstance(self.data_extractor, _ChEBIDataExtractor)
+            or self.data_extractor is None
+        )
+        super().__init__(**kwargs)
 
     def set_pos_weight(self, input: torch.Tensor) -> None:
         """
@@ -50,6 +56,9 @@ def set_pos_weight(self, input: torch.Tensor) -> None:
             )
             and self.pos_weight is None
         ):
+            print(
+                f"Computing loss-weights based on v{self.data_extractor.chebi_version} dataset (beta={self.beta})"
+            )
             complete_data = pd.concat(
                 [
                     pd.read_pickle(
@@ -75,7 +84,9 @@ def set_pos_weight(self, input: torch.Tensor) -> None:
                 [w / mean for w in weights], device=input.device
             )
 
-    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, input: torch.Tensor, target: torch.Tensor, **kwargs
+    ) -> torch.Tensor:
         """
         Forward pass for the loss calculation.