File size: 989 Bytes
590a501 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """XGBoost model wrapper."""
from __future__ import annotations
import xgboost as xgb
import pandas as pd
from sklearn.metrics import mean_squared_error
def train_xgboost(train_df: pd.DataFrame, valid_df: pd.DataFrame, feature_cols: list[str], target_col: str = "target_rank"):
dtrain = xgb.DMatrix(train_df[feature_cols], label=train_df[target_col])
dvalid = xgb.DMatrix(valid_df[feature_cols], label=valid_df[target_col])
params = {
"objective": "reg:squarederror",
"learning_rate": 0.05,
"max_depth": 6,
"subsample": 0.8,
"colsample_bytree": 0.8,
"seed": 42,
}
model = xgb.train(
params,
dtrain,
num_boost_round=500,
evals=[(dvalid, "valid")],
early_stopping_rounds=50,
verbose_eval=False,
)
pred = model.predict(dvalid)
rmse = mean_squared_error(valid_df[target_col], pred, squared=False)
print(f"XGBoost valid RMSE: {rmse:.6f}")
return model
|