Fixed initialization and formatting.

microsoft · Dec 17, 2023 · 8027c77 · 8027c77
1 parent 2fe1d2f
commit 8027c77
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 14 deletions.
diff --git a/raimitigations/automitigator/automitigator_definitions.py b/raimitigations/automitigator/automitigator_definitions.py
@@ -11,10 +11,10 @@ class AutoMitigatorDefinitions:
     synthesizer = "synthesizer"
     synthesizer_epochs_key = "epochs"
     synthesizer_model_key = "model"
-    
+
     rebalancer = "rebalancer"
     rebalancer_strategy_key = "strategy"
-    
+
     scaler = "scaler"
     scaler_name_key = "scaler_name"
     standard_scaler = "standard_scaler"
@@ -23,13 +23,13 @@ class AutoMitigatorDefinitions:
     power_scaler = "power_scaler"
     normalize_scaler = "normalize_scaler"
     minmax_scaler = "minmax_scaler"
-    
+
     imputer = "imputer"
     imputer_name_key = "imputer_name"
     basic_imputer = "basic"
     iterative_imputer = "iterative"
     knn_imputer = "knn"
-    
+
     feature_selector = "feature_selector"
     selector_name_key = "selector_name"
     sequential_selector = "sequential_selector"
@@ -42,9 +42,9 @@ class AutoMitigatorDefinitions:
     cs_test_size_key = "test_size"
     cs_algorithm_key = "algorithm"
     cs_steps_key = "steps"
-    
+
     no_mitigation = "nomitigation"
-    
+
     results_loss_key = "loss"
     results_automl_key = "automl"
-    results_pipeline_key = "pipeline"
+    results_pipeline_key = "pipeline"
diff --git a/raimitigations/automitigator/evaluator.py b/raimitigations/automitigator/evaluator.py
@@ -7,6 +7,7 @@
 from .mitigation_actions import MitigationActions
 from .automitigator_definitions import AutoMitigatorDefinitions as amd
 
+
 class Evaluator:
     """
     Evaluates a given set of mitigations on a given dataset.
@@ -47,6 +48,7 @@ def evaluate(self, train_x, train_y, search_config):
         #               {'action0': {'type': 0, 'strategy': 0, 'name': 'rebalancer'}}}}
 
         self.pipeline = None
+        self.pipeline_steps = []
         search_space = search_config[amd.search_space_key]
         cohort = search_space[amd.cohort_key]
         if cohort == amd.all_cohort:
@@ -114,11 +116,21 @@ def mitigate_full_dataset(self, train_x, train_y, search_space):
 
             if mitigation_name == amd.synthesizer:
                 self._pipeline_append(
-                    (amd.synthesizer, MitigationActions.get_synthesizer(config[amd.synthesizer_epochs_key], config[amd.synthesizer_model_key]))
+                    (
+                        amd.synthesizer,
+                        MitigationActions.get_synthesizer(
+                            config[amd.synthesizer_epochs_key], config[amd.synthesizer_model_key]
+                        ),
+                    )
                 )
             elif mitigation_name == amd.rebalancer:
                 self._pipeline_append(
-                    (amd.rebalancer, MitigationActions.get_rebalancer(config[amd.mitigation_type_key], config[amd.rebalancer_strategy_key]))
+                    (
+                        amd.rebalancer,
+                        MitigationActions.get_rebalancer(
+                            config[amd.mitigation_type_key], config[amd.rebalancer_strategy_key]
+                        ),
+                    )
                 )
             elif mitigation_name == amd.scaler:
                 self._process_scaler(config[amd.mitigation_type_key])

diff --git a/raimitigations/automitigator/searchspacebuilder.py b/raimitigations/automitigator/searchspacebuilder.py
@@ -1,12 +1,13 @@
 from flaml import tune
 from .automitigator_definitions import AutoMitigatorDefinitions as amd
 
+
 class SearchSpaceBuilder:
     """
     SearchSpaceBuilder is a class for building the search space for the AutoML
     hyperparameter search.
-    All potential mitigations are included in the search space. The search space is modeled 
-    as a hyperparameter optimization problem, where the hyperparameters are each of the mitigation's 
+    All potential mitigations are included in the search space. The search space is modeled
+    as a hyperparameter optimization problem, where the hyperparameters are each of the mitigation's
     parameters and the values are the possible values for each parameter.
 
     :param int max_mitigations: The maximum number of mitigations to be applied to the dataset.
@@ -28,13 +29,17 @@ def build(self):
         synthesizer = {
             amd.mitigation_name_key: amd.synthesizer,
             amd.synthesizer_epochs_key: tune.randint(lower=200, upper=700),
-            amd.synthesizer_model_key: tune.randint(lower=0, upper=4),  # 0 - ctgan, 1 - tvae, 2 - copula, 3 - copula_gan
+            amd.synthesizer_model_key: tune.randint(
+                lower=0, upper=4
+            ),  # 0 - ctgan, 1 - tvae, 2 - copula, 3 - copula_gan
         }
 
         rebalancer = {
             amd.mitigation_name_key: amd.rebalancer,
             amd.mitigation_type_key: tune.randint(lower=0, upper=3),  # 0 - oversample, 1 - undersample, 2 - both
-            amd.rebalancer_strategy_key: tune.randint(lower=0, upper=4),  # 0 - majority, 1 - not minority, 2 - not majority, 3 - all
+            amd.rebalancer_strategy_key: tune.randint(
+                lower=0, upper=4
+            ),  # 0 - majority, 1 - not minority, 2 - not majority, 3 - all
         }
 
         standard_scaler = {amd.scaler_name_key: amd.standard_scaler}
@@ -62,7 +67,10 @@ def build(self):
 
         knn_imputer = {amd.imputer_name_key: amd.knn_imputer}
 
-        imputer = {amd.mitigation_name_key: amd.imputer, amd.mitigation_type_key: tune.choice([basic_imputer, iterative_imputer, knn_imputer])}
+        imputer = {
+            amd.mitigation_name_key: amd.imputer,
+            amd.mitigation_type_key: tune.choice([basic_imputer, iterative_imputer, knn_imputer]),
+        }
 
         sequential_selector = {
             amd.selector_name_key: amd.sequential_selector,