@@ -571,21 +571,21 @@ def lgbm_classification_learner(df: pd.DataFrame,
571
571
params = assoc (params , "eta" , learning_rate )
572
572
params = params if "objective" in params else assoc (params , "objective" , 'binary' )
573
573
574
- weights = df [weight_column ]. values if weight_column else None
574
+ weights = df [weight_column ] if weight_column else None
575
575
576
576
features = features if not encode_extra_cols else expand_features_encoded (df , features )
577
577
578
578
dtrain = lgbm .Dataset (df [features ], label = df [target ], feature_name = list (map (str , features )), weight = weights ,
579
579
silent = True , categorical_feature = categorical_features )
580
580
581
- bst = lgbm .train (params , dtrain , num_estimators )
581
+ bst = lgbm .train (params , dtrain , num_estimators , categorical_feature = categorical_features )
582
582
583
583
def p (new_df : pd .DataFrame , apply_shap : bool = False ) -> pd .DataFrame :
584
584
if params ["objective" ] == "multiclass" :
585
585
col_dict = {prediction_column + "_" + str (key ): value
586
- for (key , value ) in enumerate (bst .predict (new_df [features ]. values ).T )}
586
+ for (key , value ) in enumerate (bst .predict (new_df [features ]).T )}
587
587
else :
588
- col_dict = {prediction_column : bst .predict (new_df [features ]. values )}
588
+ col_dict = {prediction_column : bst .predict (new_df [features ])}
589
589
590
590
if apply_shap :
591
591
import shap
0 commit comments