test for chebai cli, trains mlp

Author: aditya0by0

chebai/preprocessing/datasets/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)

tests/unit/cli/__init__.py

@@ -0,0 +1 @@
+class_path: chebai.preprocessing.datasets.mock_dm.MyLightningDataModule

tests/unit/cli/mock_dm_config.yml

@@ -0,0 +1,37 @@
+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",
+            "--model.test_metrics=configs/metrics/micro-macro-f1.yml",
+            "--model.val_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()