src/models/tune.py

import optuna
import mlflow
import mlflow.sklearn
import numpy as np
from src.utils.mlflow_utils import setup_mlflow

from sklearn.ensemble import RandomForestRegressor
from sklearn.pipeline import Pipeline
from sklearn.metrics import mean_squared_error


def tune_random_forest(X_train, X_test, y_train, y_test, preprocessor):

    setup_mlflow("AQI_Prediction")

    def objective(trial):
        with mlflow.start_run(nested=True):

            params = {
                "n_estimators": trial.suggest_int("n_estimators", 50, 300),
                "max_depth": trial.suggest_int("max_depth", 5, 30),
                "min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
                "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 5),
                "random_state": 42,
                "n_jobs": -1
            }

            mlflow.log_params(params)

            model = RandomForestRegressor(**params)

            pipeline = Pipeline([
                ("preprocessor", preprocessor),
                ("model", model)
            ])

            pipeline.fit(X_train, y_train)
            y_pred = pipeline.predict(X_test)

            rmse = mean_squared_error(y_test, y_pred) ** 0.5
            mlflow.log_metric("rmse", rmse)

            return rmse

    with mlflow.start_run(run_name="Optuna_RF_Search"):
        study = optuna.create_study(direction="minimize")
        study.optimize(objective, n_trials=50)

        mlflow.log_params(study.best_params)
        mlflow.log_metric("best_rmse", study.best_value)

    return study.best_params