Get estimator name

VIDA-NYU · Mar 16, 2024 · 75faf92 · 75faf92
1 parent 70303c8
commit 75faf92
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 15 deletions.
diff --git a/alpha_automl/automl_api.py b/alpha_automl/automl_api.py
@@ -105,21 +105,7 @@ def fit(self, X, y):
         for index, pipeline in enumerate(sorted_pipelines, start=1):
             pipeline_id = PIPELINE_PREFIX + str(index)
             self.pipelines[pipeline_id] = pipeline
-            if (
-                pipeline.get_pipeline().steps[-1][0]
-                == 'sklearn.semi_supervised.SelfTrainingClassifier'
-                or pipeline.get_pipeline().steps[-1][0]
-                == 'alpha_automl.builtin_primitives.semisupervised_classifier.AutonBox'
-            ):
-                leaderboard_data.append(
-                    [
-                        index,
-                        f'{pipeline.get_summary()}, {pipeline.get_pipeline().steps[-1][1].base_estimator.__class__.__name__}',
-                        pipeline.get_score(),
-                    ]
-                )
-            else:
-                leaderboard_data.append([index, pipeline.get_summary(), pipeline.get_score()])
+            leaderboard_data.append([index, pipeline.get_summary(), pipeline.get_score()])
 
         self.leaderboard = pd.DataFrame(leaderboard_data, columns=['ranking', 'pipeline', self.metric])
 

diff --git a/alpha_automl/pipeline.py b/alpha_automl/pipeline.py
@@ -1,5 +1,7 @@
 import logging
 from sklearn.compose import ColumnTransformer
+from sklearn.semi_supervised import SelfTrainingClassifier
+from alpha_automl.builtin_primitives.semisupervised_classifier import AutonBox
 from alpha_automl.utils import COLUMN_SELECTOR_ID
 
 
@@ -56,5 +58,8 @@ def _make_summary(self):
                     step_name = transformer_name.split('-')[0].split('.')[-1]
                     if step_name not in step_names:
                         step_names.append(step_name)
+            elif isinstance(step_object, SelfTrainingClassifier) or isinstance(step_object, AutonBox):
+                estimator_name = step_object.base_estimator.__class__.__name__
+                step_names.append(estimator_name)
 
         self.summary = ', '.join(step_names)