Hyperparameter Tuning with MLflow

import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import RandomizedSearchCV
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from scipy.stats import randint, uniform

wine = load_wine()
X_train, X_test, y_train, y_test = train_test_split(
    wine.data, wine.target, test_size=0.2, random_state=42
)

param_distributions = {
    "n_estimators": randint(50, 500),
    "max_depth": randint(3, 20),
    "min_samples_split": randint(2, 20),
    "min_samples_leaf": randint(1, 10),
    "max_features": uniform(0.1, 0.9),
}

mlflow.autolog()

with mlflow.start_run(run_name="hyperparameter-search"):
    search = RandomizedSearchCV(
        RandomForestClassifier(random_state=42),
        param_distributions=param_distributions,
        n_iter=50,
        cv=5,
        scoring="accuracy",
        random_state=42,
        n_jobs=-1,
    )
    search.fit(X_train, y_train)

    # Log best results
    mlflow.log_metric("best_cv_score", search.best_score_)
    mlflow.log_params(
        {f"best_{k}": v for k, v in search.best_params_.items()}
    )

    # Evaluate on test set
    test_accuracy = search.score(X_test, y_test)
    mlflow.log_metric("test_accuracy", test_accuracy)

    print(f"Best CV Score: {search.best_score_:.4f}")
    print(f"Test Accuracy: {test_accuracy:.4f}")
    print(f"Best Params: {search.best_params_}")