fix collated labels none error

Author: aditya0by0

chebai/ensemble/__init__.py

@@ -1,5 +1,7 @@
+from ._base import EnsembleBase
 from ._consolidator import WeightedMajorityVoting
 from ._controller import NoActivationCondition
+from ._wrappers import NNWrapper
 
 
 class FullEnsembleWMV(NoActivationCondition, WeightedMajorityVoting):
@@ -8,4 +10,4 @@ class FullEnsembleWMV(NoActivationCondition, WeightedMajorityVoting):
     pass
 
 
-__all__ = ["FullEnsembleWMV"]
+__all__ = ["FullEnsembleWMV", "NNWrapper"]

chebai/ensemble/_base.py

@@ -205,6 +205,9 @@ def run_ensemble(self) -> None:
         )
 
         if self._operation == EVAL_OP:
+            assert (
+                self._collated_labels is not None
+            ), "Collated labels must be set for evaluation operation."
             print_metrics(
                 final_preds,
                 self._collated_labels,

chebai/ensemble/_controller.py

@@ -7,7 +7,7 @@
 from torch import Tensor
 
 from ._base import EnsembleBase
-from ._constants import EVAL_OP, PRED_OP, WRAPPER_CLS_PATH
+from ._constants import PRED_OP, WRAPPER_CLS_PATH
 from ._utils import load_class
 from ._wrappers import BaseWrapper
 
@@ -35,7 +35,7 @@ def __init__(self, **kwargs: Any):
         # This is in order to avoid re-adding models that have already been processed
         self._model_key_set: set[str] = set(self._model_configs.keys())
 
-        # Labels from any processed data.pt file for any reader
+        # Labels from any processed `data.pt` file of any reader
         self._collated_labels: torch.Tensor | None = None
 
     def _controller(
@@ -56,6 +56,12 @@ def _controller(
             model_output, model_props = wrapped_model.predict(model_input)
         else:
             model_output, model_props = wrapped_model.evaluate(model_input)
+            if (
+                self._collated_labels is None
+                and wrapped_model.collated_labels is not None
+            ):
+                self._collated_labels = wrapped_model.collated_labels
+
         del wrapped_model  # Model can be huge to keep it in memory, delete asap as no longer needed
 
         pred_conf_dict = self._get_pred_conf_from_model_output(
@@ -98,8 +104,6 @@ def _wrap_model(self, model_name: str) -> BaseWrapper:
             dm_labels=self._dm_labels,
             **self._kwargs
         )
-        if self._collated_labels is not None and self._operation == EVAL_OP:
-            self._collated_labels = wrapped_model.collated_labels
 
         assert isinstance(wrapped_model, BaseWrapper), ""
         return wrapped_model

chebai/ensemble/_scripts/_ensemble_run_script.py

@@ -1,11 +1,10 @@
-from typing import Any, Dict, Type
+from typing import Any, Dict
 
 import yaml
 from jsonargparse import ArgumentParser
 
-from chebai.ensemble._utils import _load_class
-
-from .._base import EnsembleBase
+from chebai.ensemble._base import EnsembleBase
+from chebai.ensemble._utils import load_class
 
 
 def load_config_and_instantiate(config_path: str) -> EnsembleBase:
@@ -27,7 +26,7 @@ def load_config_and_instantiate(config_path: str) -> EnsembleBase:
     class_path: str = config["class_path"]
     init_args: Dict[str, Any] = config.get("init_args", {})
 
-    cls = _load_class(class_path)
+    cls = load_class(class_path)
 
     if not issubclass(cls, EnsembleBase):
         raise TypeError(f"{cls} must be subclass of EnsembleBase")