Fix a bug when target class is a string, but cat encoing via ml algor…

…ithm is chosen
ThomasMeissnerDS · Aug 1, 2024 · d7f161d · d7f161d
1 parent 3af3ed5
commit d7f161d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/bluecast/blueprints/cast.py b/bluecast/blueprints/cast.py
@@ -397,8 +397,9 @@ def fit(self, df: pd.DataFrame, target_col: str) -> None:
             )
             x_test = self.cat_encoder.transform_target_encode_multiclass(x_test.copy())
         elif self.conf_training.cat_encoding_via_ml_algorithm:
-            x_train[self.cat_columns] = x_train[self.cat_columns].astype("category")
-            x_test[self.cat_columns] = x_test[self.cat_columns].astype("category")
+            cat_cols = [col for col in self.cat_columns if col != self.target_column]
+            x_train[cat_cols] = x_train[cat_cols].astype("category")
+            x_test[cat_cols] = x_test[cat_cols].astype("category")
 
         if self.custom_last_mile_computation:
             x_train, y_train = self.custom_last_mile_computation.fit_transform(
@@ -587,7 +588,8 @@ def transform_new_data(self, df: pd.DataFrame) -> pd.DataFrame:
         ):
             df = self.cat_encoder.transform_target_encode_multiclass(df.copy())
         elif self.conf_training.cat_encoding_via_ml_algorithm:
-            df[self.cat_columns] = df[self.cat_columns].astype("category")
+            cat_cols = [col for col in self.cat_columns if col != self.target_column]
+            df[cat_cols] = df[cat_cols].astype("category")
 
         if self.custom_last_mile_computation:
             df, _ = self.custom_last_mile_computation.transform(

diff --git a/bluecast/blueprints/cast_regression.py b/bluecast/blueprints/cast_regression.py
@@ -375,8 +375,9 @@ def fit(self, df: pd.DataFrame, target_col: str) -> None:
                 x_test.copy()
             )
         elif self.conf_training.cat_encoding_via_ml_algorithm:
-            x_train[self.cat_columns] = x_train[self.cat_columns].astype("category")
-            x_test[self.cat_columns] = x_test[self.cat_columns].astype("category")
+            cat_cols = [col for col in self.cat_columns if col != self.target_column]
+            x_train[cat_cols] = x_train[cat_cols].astype("category")
+            x_test[cat_cols] = x_test[cat_cols].astype("category")
 
         if self.custom_last_mile_computation:
             x_train, y_train = self.custom_last_mile_computation.fit_transform(
@@ -544,7 +545,8 @@ def transform_new_data(self, df: pd.DataFrame) -> pd.DataFrame:
         ):
             df = self.cat_encoder.transform_target_encode_binary_class(df.copy())
         elif self.conf_training.cat_encoding_via_ml_algorithm:
-            df[self.cat_columns] = df[self.cat_columns].astype("category")
+            cat_cols = [col for col in self.cat_columns if col != self.target_column]
+            df[cat_cols] = df[cat_cols].astype("category")
 
         if self.custom_last_mile_computation:
             df, _ = self.custom_last_mile_computation.transform(