ChEB-AI · sfluegel05 · Sep 16, 2025 · Jul 15, 2025 · Jul 15, 2025 · Jul 15, 2025
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -9,7 +9,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12"]
 
     steps:
       - uses: actions/checkout@v4
@@ -24,7 +24,8 @@ jobs:
           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 .
+          # wandb included to use custom trainer for cli test which needs wandb logger
+          python -m pip install -e .[wandb]
 
       - name: Display Python & Installed Packages
         run: |

diff --git a/chebai/models/ffn.py b/chebai/models/ffn.py
@@ -11,7 +11,6 @@ class FFN(ChebaiBaseNet):
 
     def __init__(
         self,
-        input_size: int,
         hidden_layers: List[int] = [
             1024,
         ],
@@ -20,7 +19,7 @@ def __init__(
         super().__init__(**kwargs)
 
         layers = []
-        current_layer_input_size = input_size
+        current_layer_input_size = self.input_dim
         for hidden_dim in hidden_layers:
             layers.append(MLPBlock(current_layer_input_size, hidden_dim))
             layers.append(Residual(MLPBlock(hidden_dim, hidden_dim)))

diff --git a/configs/model/ffn.yml b/configs/model/ffn.yml
@@ -2,4 +2,3 @@ class_path: chebai.models.ffn.FFN
 init_args:
   optimizer_kwargs:
     lr: 1e-3
-  input_size: 2560
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ maintainers = [
 ]
 readme = "README.md"
 license = { text = "AGPL-3.0" }
-requires-python = ">=3.9,<3.13"
+requires-python = ">=3.10,<3.13"
 dependencies = [
     "networkx",
     "numpy",
@@ -28,7 +28,7 @@ dependencies = [
     "pysmiles==1.1.2",
     "rdkit",
     "selfies",
-    "lightning>=2.5",
+    "lightning<=2.5.1",
     "jsonargparse[signatures]>=4.17",
     "omegaconf",
     "deepsmiles",

diff --git a/tests/unit/cli/__init__.py b/tests/unit/cli/__init__.py
diff --git a/tests/unit/cli/mock_dm.py b/tests/unit/cli/mock_dm.py
@@ -0,0 +1,46 @@
+import torch
+from lightning.pytorch.core.datamodule import LightningDataModule
+from torch.utils.data import DataLoader
+
+from chebai.preprocessing.collate import RaggedCollator
+
+
+class MyLightningDataModule(LightningDataModule):
+    def __init__(self):
+        super().__init__()
+        self._num_of_labels = None
+        self._feature_vector_size = None
+        self.collator = RaggedCollator()
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None):
+        self._num_of_labels = 10
+        self._feature_vector_size = 20
+        print(f"Number of labels: {self._num_of_labels}")
+        print(f"Number of features: {self._feature_vector_size}")
+
+    @property
+    def num_of_labels(self):
+        return self._num_of_labels
+
+    @property
+    def feature_vector_size(self):
+        return self._feature_vector_size
+
+    def train_dataloader(self):
+        assert self.feature_vector_size is not None, "feature_vector_size must be set"
+        # Dummy dataset for example purposes
+
+        datalist = [
+            {
+                "features": torch.randn(self._feature_vector_size),
+                "labels": torch.randint(0, 2, (self._num_of_labels,), dtype=torch.bool),
+                "ident": i,
+                "group": None,
+            }
+            for i in range(100)
+        ]
+
+        return DataLoader(datalist, batch_size=32, collate_fn=self.collator)
diff --git a/tests/unit/cli/mock_dm_config.yml b/tests/unit/cli/mock_dm_config.yml
@@ -0,0 +1 @@
+class_path: tests.unit.cli.mock_dm.MyLightningDataModule
diff --git a/tests/unit/cli/testCLI.py b/tests/unit/cli/testCLI.py
@@ -0,0 +1,35 @@
+import unittest
+
+from chebai.cli import ChebaiCLI
+
+
+class TestChebaiCLI(unittest.TestCase):
+    def setUp(self):
+        self.cli_args = [
+            "fit",
+            "--trainer=configs/training/default_trainer.yml",
+            "--model=configs/model/ffn.yml",
+            "--model.init_args.hidden_layers=[10]",
+            "--model.train_metrics=configs/metrics/micro-macro-f1.yml",
+            "--data=tests/unit/cli/mock_dm_config.yml",
+            "--model.pass_loss_kwargs=false",
+            "--trainer.min_epochs=1",
+            "--trainer.max_epochs=1",
+            "--model.criterion=configs/loss/bce.yml",
+            "--model.criterion.init_args.beta=0.99",
+        ]
+
+    def test_mlp_on_chebai_cli(self):
+        # Instantiate ChebaiCLI and ensure no exceptions are raised
+        try:
+            ChebaiCLI(
+                args=self.cli_args,
+                save_config_kwargs={"config_filename": "lightning_config.yaml"},
+                parser_kwargs={"parser_mode": "omegaconf"},
+            )
+        except Exception as e:
+            self.fail(f"ChebaiCLI raised an unexpected exception: {e}")
+
+
+if __name__ == "__main__":
+    unittest.main()
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		class_path: tests.unit.cli.mock_dm.MyLightningDataModule