Fit a bayesian optimizer and a decsion tree model

decison_tree_and_bayes_optimizer.py

# Define a model 
tree = DecisionTreeRegressor(criterion="mse")

# Create the Bayesion optimization object
opt = BayesSearchCV(
    tree,
    {
        "max_depth": (5, 15),
        "splitter": ["best", "random"],
    },
    n_iter=n_iter,
    cv=n_folds
)

# Train the model
opt.fit(train_x, train_y)

model = opt.best_estimator_

# Save the trained model
joblib.dump(model, os.path.join(args.model_dir, "model.joblib"))