models/xgboost/config.py

"""Shared XGBoost config used by training and evaluation. Loads from config/default.yaml when present."""

from copy import deepcopy

from xgboost import XGBClassifier


def _load_xgb_params():
    try:
        from config import get
        xgb = get("xgboost") or {}
        return {
            "n_estimators": xgb.get("n_estimators", 600),
            "max_depth": xgb.get("max_depth", 8),
            "learning_rate": xgb.get("learning_rate", 0.1489),
            "subsample": xgb.get("subsample", 0.9625),
            "colsample_bytree": xgb.get("colsample_bytree", 0.9013),
            "reg_alpha": xgb.get("reg_alpha", 1.1407),
            "reg_lambda": xgb.get("reg_lambda", 2.4181),
            "eval_metric": xgb.get("eval_metric", "logloss"),
        }
    except Exception:
        return {
            "n_estimators": 600,
            "max_depth": 8,
            "learning_rate": 0.1489,
            "subsample": 0.9625,
            "colsample_bytree": 0.9013,
            "reg_alpha": 1.1407,
            "reg_lambda": 2.4181,
            "eval_metric": "logloss",
        }


XGB_BASE_PARAMS = _load_xgb_params()


def get_xgb_params():
    return deepcopy(XGB_BASE_PARAMS)


def build_xgb_classifier(seed: int, *, verbosity: int = 0, early_stopping_rounds=None):
    params = get_xgb_params()
    params.update(
        {
            "random_state": seed,
            "verbosity": verbosity,
        }
    )
    if early_stopping_rounds is not None:
        params["early_stopping_rounds"] = early_stopping_rounds
    return XGBClassifier(**params)