mirror of
https://github.com/webclinic017/drift.git
synced 2026-08-16 04:18:07 +00:00
feat(Models): added optional hpsklearn (#236)
This commit is contained in:
+14
-1
@@ -1,5 +1,5 @@
|
||||
from models.sklearn import SKLearnModel
|
||||
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
|
||||
from sklearn.ensemble import RandomForestClassifier
|
||||
from .base import Model
|
||||
|
||||
default_feature_selector_classification = SKLearnModel(
|
||||
@@ -75,6 +75,19 @@ def get_model(model_name: str) -> Model:
|
||||
SKLearnModel(LGBMClassifier(n_jobs=-1, max_depth=20, random_state=1))
|
||||
)
|
||||
|
||||
elif model_name == "HyperOpt":
|
||||
from hpsklearn import HyperoptEstimator
|
||||
from hyperopt import tpe
|
||||
|
||||
return set_name(
|
||||
SKLearnModel(
|
||||
HyperoptEstimator(
|
||||
algo=tpe.suggest,
|
||||
trial_timeout=300,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
elif model_name == "AutoML":
|
||||
from supervised.automl import AutoML
|
||||
|
||||
|
||||
Reference in New Issue
Block a user