models/xgboost_model.py

"""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