personalize encoders in json

Author: diegomarvid

pipeline_lib/core/steps/encode.py

@@ -1,4 +1,5 @@
-from typing import List, Optional, Tuple
+import json
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -16,11 +17,20 @@ class EncodeStep(PipelineStep):
     used_for_prediction = True
     used_for_training = True
 
-    def __init__(self, target: Optional[str] = None, cardinality_threshold: float = 0.3) -> None:
+    def __init__(
+        self,
+        target: Optional[str] = None,
+        cardinality_threshold: float = 0.3,
+        low_cardinality_encoder: str = "OrdinalEncoder",
+        high_cardinality_encoder: str = "TargetEncoder",
+    ) -> None:
         """Initialize EncodeStep."""
         self.init_logger()
         self.target = target
         self.cardinality_threshold = cardinality_threshold
+        self.low_cardinality_encoder = low_cardinality_encoder
+        self.high_cardinality_encoder = high_cardinality_encoder
+        self.encoder_feature_map = {}
 
     def execute(self, data: DataContainer) -> DataContainer:
         """Execute the encoding step."""
@@ -39,28 +49,26 @@ def execute(self, data: DataContainer) -> DataContainer:
         if pd.api.types.is_numeric_dtype(df[target_column_name]):
             target_original_dtype = df[target_column_name].dtype
 
-        self._log_feature_info(
-            categorical_features,
-            numeric_features,
-            low_cardinality_features,
-            high_cardinality_features,
-        )
-
         column_transformer = self._create_column_transformer(
             high_cardinality_features, low_cardinality_features
         )
 
         encoded_data = self._transform_data(df, target_column_name, column_transformer)
         encoded_data = self._restore_column_order(df, encoded_data)
-        encoded_data = self._convert_ordinal_encoded_columns_to_int(
-            encoded_data, column_transformer
-        )
+        encoded_data = self._convert_ordinal_encoded_columns_to_int(encoded_data)
         encoded_data = self._restore_numeric_dtypes(encoded_data, original_numeric_dtypes)
         encoded_data = self._restore_target_dtype(
             encoded_data, target_column_name, target_original_dtype
         )
         encoded_data = self._convert_float64_to_float32(encoded_data)
 
+        self._log_feature_info(
+            categorical_features,
+            numeric_features,
+            low_cardinality_features,
+            high_cardinality_features,
+        )
+
         data.flow = encoded_data
 
         return data
@@ -93,14 +101,50 @@ def _split_categorical_features(
         ]
         return low_cardinality_features, high_cardinality_features
 
+    def _get_encoder(self, encoder_name: str) -> Union[OrdinalEncoder, TargetEncoder]:
+        """Map encoder name to the corresponding encoder class."""
+        encoder_map = {
+            "OrdinalEncoder": OrdinalEncoder(),
+            "TargetEncoder": TargetEncoder(),
+            # Add more encoders as needed
+        }
+
+        encoder = encoder_map.get(encoder_name)
+
+        if not encoder:
+            raise ValueError(
+                f"Unsupported encoder: {encoder_name}. Supported encoders: {encoder_map}"
+            )
+
+        return encoder
+
     def _create_column_transformer(
         self, high_cardinality_features: List[str], low_cardinality_features: List[str]
     ) -> ColumnTransformer:
         """Create a ColumnTransformer for encoding."""
+        high_cardinality_encoder = self._get_encoder(self.high_cardinality_encoder)
+        low_cardinality_encoder = self._get_encoder(self.low_cardinality_encoder)
+
+        # Initialize the encoder_feature_map as an empty dictionary
+        self.encoder_feature_map = {}
+
+        # Check if both encoders are the same
+        if self.high_cardinality_encoder == self.low_cardinality_encoder:
+            # If the same, merge the feature lists
+            # This assumes you want to combine the features into a single list; adjust if needed
+            combined_features = high_cardinality_features + low_cardinality_features
+            self.encoder_feature_map[self.high_cardinality_encoder] = combined_features
+        else:
+            # If not the same, assign individually
+            self.encoder_feature_map[self.high_cardinality_encoder] = high_cardinality_features
+            self.encoder_feature_map[self.low_cardinality_encoder] = low_cardinality_features
+
+        print(self.encoder_feature_map)
+
         return ColumnTransformer(
             [
-                ("target_encoder", TargetEncoder(), high_cardinality_features),
-                ("ordinal_encoder", OrdinalEncoder(), low_cardinality_features),
+                ("high_cardinality_encoder", high_cardinality_encoder, high_cardinality_features),
+                ("low_cardinality_encoder", low_cardinality_encoder, low_cardinality_features),
             ],
             remainder="passthrough",
             verbose_feature_names_out=False,
@@ -120,15 +164,11 @@ def _restore_column_order(self, df: pd.DataFrame, encoded_data: pd.DataFrame) ->
         new_column_order = [col for col in df.columns if col in encoded_data.columns]
         return encoded_data[new_column_order]
 
-    def _convert_ordinal_encoded_columns_to_int(
-        self, encoded_data: pd.DataFrame, column_transformer: ColumnTransformer
-    ) -> pd.DataFrame:
+    def _convert_ordinal_encoded_columns_to_int(self, encoded_data: pd.DataFrame) -> pd.DataFrame:
         """Convert ordinal encoded columns to the smallest possible integer dtype."""
-        ordinal_encoder_features = column_transformer.named_transformers_[
-            "ordinal_encoder"
-        ].get_feature_names_out()
+        ordinal_encoded_features = self.encoder_feature_map.get("OrdinalEncoder", [])
 
-        for col in ordinal_encoder_features:
+        for col in ordinal_encoded_features:
             if col in encoded_data.columns:
                 n_unique = encoded_data[col].nunique()
                 if n_unique <= 2**8:
@@ -199,11 +239,9 @@ def _log_feature_info(
             f"Low cardinality features (cardinality ratio < {self.cardinality_threshold}):"
             f" ({len(low_cardinality_features)}) - {low_cardinality_features}"
         )
-        self.logger.info("Low cardinality features encoding method: ordinal encoder")
         self.logger.info(
             f"High cardinality features (cardinality ratio >= {self.cardinality_threshold}):"
             f" ({len(high_cardinality_features)}) -  {high_cardinality_features}"
         )
-        self.logger.info("High cardinality features encoding method: target encoder")
         self.logger.info(f"Numeric features: ({len(numeric_features)}) - {numeric_features}")
-        self.logger.info("Numeric features encoding method: passthrough")
+        self.logger.info(f"Encoder feature map: \n{json.dumps(self.encoder_feature_map, indent=4)}")