From 63961f4bb680bffac7b8e5dde9d48e4b518d0be9 Mon Sep 17 00:00:00 2001 From: Fabio Beranizo Date: Thu, 19 May 2022 21:58:13 -0300 Subject: [PATCH] Fix typing annotation of categorical_features It is a Union[List[str], 'auto'] --- src/fklearn/training/classification.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/fklearn/training/classification.py b/src/fklearn/training/classification.py index 6cb4784f..3e30b377 100644 --- 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: