lightning module: retrieve num_of_labels from data module

Author: aditya0by0

chebai/cli.py

@@ -38,15 +38,15 @@ def add_arguments_to_parser(self, parser: LightningArgumentParser):
         Args:
             parser (LightningArgumentParser): Argument parser instance.
         """
-        for kind in ("train", "val", "test"):
-            for average in ("micro-f1", "macro-f1", "balanced-accuracy"):
-                parser.link_arguments(
-                    "model.init_args.out_dim",
-                    f"model.init_args.{kind}_metrics.init_args.metrics.{average}.init_args.num_labels",
-                )
-        parser.link_arguments(
-            "model.init_args.out_dim", "trainer.callbacks.init_args.num_labels"
-        )
+        # for kind in ("train", "val", "test"):
+        #     for average in ("micro-f1", "macro-f1", "balanced-accuracy"):
+        #         parser.link_arguments(
+        #             "model.init_args.out_dim",
+        #             f"model.init_args.{kind}_metrics.init_args.metrics.{average}.init_args.num_labels",
+        #         )
+        # parser.link_arguments(
+        #     "model.init_args.out_dim", "trainer.callbacks.init_args.num_labels"
+        # )
         parser.link_arguments(
             "data", "model.init_args.criterion.init_args.data_extractor"
         )

chebai/models/base.py

@@ -35,7 +35,6 @@ class ChebaiBaseNet(LightningModule):
     def __init__(
         self,
         criterion: torch.nn.Module = None,
-        out_dim: Optional[int] = None,
         train_metrics: Optional[torch.nn.Module] = None,
         val_metrics: Optional[torch.nn.Module] = None,
         test_metrics: Optional[torch.nn.Module] = None,
@@ -48,7 +47,7 @@ def __init__(
         self.save_hyperparameters(
             ignore=["criterion", "train_metrics", "val_metrics", "test_metrics"]
         )
-        self.out_dim = out_dim
+        self.out_dim = None
         if optimizer_kwargs:
             self.optimizer_kwargs = optimizer_kwargs
         else:
@@ -70,6 +69,14 @@ def __init_subclass__(cls, **kwargs):
         else:
             _MODEL_REGISTRY[cls.NAME] = cls
 
+    def setup(self, stage: str) -> None:
+        if self.trainer and hasattr(self.trainer, "datamodule"):
+            self.out_dim = int(self.trainer.datamodule.hparams.num_of_labels)
+        else:
+            raise ValueError("Trainer has no data module")
+        assert self.out_dim is not None, "Model output dimension is None"
+        print(f"Output Dimension for the model: {self.out_dim}")
+
     def _get_prediction_and_labels(
         self, data: Dict[str, Any], labels: torch.Tensor, output: torch.Tensor
     ) -> (torch.Tensor, torch.Tensor):