xgboost_sklearn_wrapper_regression_quickstart_diabetes_dataset.py

1import xgboost as xgb
2from sklearn.datasets import load_diabetes
3from sklearn.model_selection import train_test_split
4from sklearn.metrics import mean_squared_error
5
6# Load the data
7X, y = load_diabetes(return_X_y=True)
8
9# Split the data into training and testing sets
10X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
11
12# Create an XGBoost regression model
13# For classification, use xgb.XGBClassifier()
14regressor = xgb.XGBRegressor(
15    n_estimators=100,
16    learning_rate=0.1,
17    max_depth=5,
18    random_state=42
19)
20
21# Fit the model
22regressor.fit(X_train, y_train)
23
24# Make predictions
25predictions = regressor.predict(X_test)
26
27# Evaluate the model
28mse = mean_squared_error(y_test, predictions)
29print(f"Mean Squared Error: {mse}")