Commit c028e5d1 authored by Leodegario Lorenzo II's avatar Leodegario Lorenzo II
Browse files

Fix advance GBM import behavior

parent 21c90869
......@@ -168,10 +168,12 @@ class MLModels:
'Multinomial NB': MNB(C),
'Complement NB': CNB(C),
'Bernoulli NB': BNB(C),
'XGB Classifier': XGBC(max_depth, tree_rs),
'LightGBM Classifier': LGBMC(max_depth, tree_rs),
'CatBoost Classifier': CBClassifier(max_depth, tree_rs),
'MLP Classifier': MLPC(C)}
if adv_gbm:
methods['Classification'].update({
'XGB Classifier': XGBC(max_depth, tree_rs),
'LightGBM Classifier': LGBMC(max_depth, tree_rs),
'CatBoost Classifier': CBClassifier(max_depth, tree_rs),})
# Set Regression methods
methods['Regression'] = {
......@@ -191,10 +193,12 @@ class MLModels:
'GB Regressor': GBRegressor(max_depth, tree_rs),
'Histogram GBR': HistGBR(max_depth, tree_rs),
'AdaBoost DT': ABDTRegressor(max_depth, tree_rs),
'XGB Regressor': XGBR(max_depth, tree_rs),
'LightGBM Regressor': LGBMR(max_depth, tree_rs),
'CatBoost Regressor': CBRegressor(max_depth, tree_rs),
'MLP Regressor': MLPR(C)}
if adv_gbm:
methods['Regression'].update({
'XGB Regressor': XGBR(max_depth, tree_rs),
'LightGBM Regressor': LGBMR(max_depth, tree_rs),
'CatBoost Regressor': CBRegressor(max_depth, tree_rs),})
return methods
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment