Fix typing annotation of categorical_features

It is a Union[List[str], 'auto']
nubank · May 20, 2022 · 63961f4 · 63961f4
1 parent c3f221f
commit 63961f4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/fklearn/training/classification.py b/src/fklearn/training/classification.py
@@ -1,4 +1,4 @@
-from typing import List, Any
+from typing import List, Any, Union
 
 import numpy as np
 import pandas as pd
@@ -502,7 +502,7 @@ def lgbm_classification_learner(df: pd.DataFrame,
                                 learning_rate: float = 0.1,
                                 num_estimators: int = 100,
                                 extra_params: LogType = None,
-                                categorical_features: List[str] = "auto",
+                                categorical_features: Union[List[str], "auto"] = "auto",
                                 prediction_column: str = "prediction",
                                 weight_column: str = None,
                                 encode_extra_cols: bool = True) -> LearnerReturnType: