tmp

VIDA-NYU · Jun 2, 2024 · 8c3ae48 · 8c3ae48
1 parent 91e9ef9
commit 8c3ae48
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/alpha_automl/hyperparameter_tuning/smac.py b/alpha_automl/hyperparameter_tuning/smac.py
@@ -99,7 +99,12 @@ def gen_configspace(pipeline):
     configspace = ConfigurationSpace(seed=0)
     all_params = {}
     for primitive, prim_obj in pipeline.steps:
-        step_type = PRIMITIVE_TYPES[primitive]
+        if "feature_engine.creation" in primitive:
+            step_type = "FEATURE_GENERATOR"
+        elif "feature_engine.selection" in primitive:
+            step_type = "FEATURE_SELECTOR"
+        else:
+            step_type = PRIMITIVE_TYPES[primitive]
         try:
             params = SMAC_DICT[primitive]
             add_params(params, all_params)