304 lines
10 KiB
Python
304 lines
10 KiB
Python
"""
|
|
Wrapper and utilities for XGBoost-based time series models integrated with MLflow.
|
|
|
|
Args:
|
|
- None
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
|
|
from typing import Any, TypeAlias, cast
|
|
|
|
import pandas as pd
|
|
import xgboost as xgb
|
|
from sientia_model.wrappers.sientia_model import SientiaModel
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
DMatrixTuple: TypeAlias = tuple[
|
|
xgb.DMatrix,
|
|
xgb.DMatrix | None,
|
|
]
|
|
|
|
|
|
class XGBoostWrapper(SientiaModel):
|
|
"""
|
|
Wrapper for XGBoost models compatible with the SientiaModel interface.
|
|
|
|
This class orchestrates preprocessing, training, retraining and prediction
|
|
using XGBoost boosters and a configured time series preprocessor so that
|
|
models can be saved and served through MLflow's pyfunc interface.
|
|
|
|
Args:
|
|
- model_type (str): Logical model type used for logging and registry.
|
|
- model_version (str): Model version identifier.
|
|
- logger (Any): Logger instance used for structured logging.
|
|
- opt_params (dict[str, Any]): Optional configuration used to tune the
|
|
underlying XGBoost booster and runtime behaviour.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
|
|
def _format_data(
|
|
self,
|
|
x: pd.DataFrame,
|
|
y: pd.Series | None = None,
|
|
x_val: pd.DataFrame | None = None,
|
|
y_val: pd.Series | None = None,
|
|
) -> DMatrixTuple:
|
|
"""
|
|
Convert pandas data into XGBoost DMatrix objects for training and validation.
|
|
|
|
Args:
|
|
- x (pd.DataFrame): Feature matrix (including any exogenous features).
|
|
- y (pd.Series | None): Optional target time series aligned with `x`.
|
|
- x_val (pd.DataFrame | None): Optional validation feature matrix.
|
|
- y_val (pd.Series | None): Optional validation target series aligned with `x_val`.
|
|
|
|
Return:
|
|
tuple[xgb.DMatrix, xgb.DMatrix | None]: DMatrix for training and an optional
|
|
DMatrix for validation when both `x_val` and `y_val` are provided.
|
|
"""
|
|
dtrain = None
|
|
dvalid = None
|
|
|
|
# Align target with features before potential index changes
|
|
if y is not None:
|
|
y = y.loc[x.index]
|
|
|
|
if 'timestamp' in x.columns:
|
|
x = x.copy()
|
|
x['timestamp'] = pd.to_datetime(x['timestamp'])
|
|
x = x.set_index('timestamp')
|
|
if y is not None:
|
|
y.index = x.index
|
|
|
|
if y is not None:
|
|
dtrain = xgb.DMatrix(x, label=y)
|
|
else:
|
|
dtrain = xgb.DMatrix(x)
|
|
|
|
if x_val is not None:
|
|
if y_val is None:
|
|
raise ValueError('When x_val is not None, y_val must be provided')
|
|
|
|
# Align validation target with validation features before index changes
|
|
y_val = y_val.loc[x_val.index]
|
|
|
|
if 'timestamp' in x_val.columns:
|
|
x_val = x_val.copy()
|
|
x_val['timestamp'] = pd.to_datetime(x_val['timestamp'])
|
|
x_val = x_val.set_index('timestamp')
|
|
y_val.index = x_val.index
|
|
|
|
dvalid = xgb.DMatrix(x_val, label=y_val)
|
|
|
|
return dtrain, dvalid
|
|
|
|
def _predict(self, data: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, Any]]:
|
|
"""
|
|
Run inference using the trained XGBoost model.
|
|
|
|
Args:
|
|
- data (pd.DataFrame): Input features for prediction; may contain a timestamp
|
|
column which is handled by `_format_data`.
|
|
|
|
Return:
|
|
tuple[pd.DataFrame, dict[str, Any]]: DataFrame with one prediction column
|
|
named after the target and an empty metadata dictionary.
|
|
"""
|
|
if self.target is None:
|
|
raise ValueError('Target column not set. Call train first.')
|
|
|
|
if not self.model_is_fitted:
|
|
raise ValueError('Model not trained. Call train first.')
|
|
|
|
X = data
|
|
|
|
dtest, _ = self._format_data(X)
|
|
predictions = self.model.predict(dtest)
|
|
|
|
return pd.DataFrame(predictions, index=X.index, columns=[self.target]), {}
|
|
|
|
def _transform(self, data: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, Any]]:
|
|
"""
|
|
Transform the raw input data using the configured time series preprocessor.
|
|
|
|
Args:
|
|
- data (pd.DataFrame): Raw input data to be transformed.
|
|
|
|
Return:
|
|
tuple[pd.DataFrame, dict[str, Any]]: Transformed feature matrix and an
|
|
empty metadata dictionary.
|
|
"""
|
|
return self.transformer.transform(data), {}
|
|
|
|
def _fit_xgboost_model(
|
|
self,
|
|
x_train: pd.DataFrame,
|
|
y_train: pd.DataFrame,
|
|
x_val: pd.DataFrame | None = None,
|
|
y_val: pd.DataFrame | None = None,
|
|
) -> None:
|
|
"""
|
|
Fit the XGBoost model with the given training and optional validation data.
|
|
|
|
Args:
|
|
- x_train (pd.DataFrame): Training feature matrix.
|
|
- y_train (pd.DataFrame): Training target data as a one-column DataFrame.
|
|
- x_val (pd.DataFrame | None): Optional validation feature matrix.
|
|
- y_val (pd.DataFrame | None): Optional validation target data.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
# Convert y to Series (squeeze can return DataFrame)
|
|
y_train_series = cast(pd.Series, y_train.squeeze())
|
|
y_val_series = cast(pd.Series, y_val.squeeze()) if y_val is not None else None
|
|
|
|
# Preprocess data using parent class method (alignment happens inside _format_data)
|
|
dtrain, dvalid = self._format_data(x_train, y_train_series, x_val, y_val_series)
|
|
params = {
|
|
'tree_method': self.opt_params.get('tree_method', 'hist'),
|
|
'device': self.opt_params.get('device', 'cuda'),
|
|
'learning_rate': self.opt_params.get('learning_rate', 0.3),
|
|
'n_estimators': self.opt_params.get('n_estimators', 100),
|
|
'max_depth': self.opt_params.get('max_depth', 32),
|
|
'subsample': self.opt_params.get('subsample', 0.8),
|
|
'colsample_bytree': self.opt_params.get('colsample_bytree', 0.8),
|
|
'min_child_weight': self.opt_params.get('min_child_weight', 5),
|
|
'random_state': self.opt_params.get('random_state', 42),
|
|
}
|
|
|
|
early_stopping_rounds: int | None = self.opt_params.get('early_stopping_rounds', None)
|
|
n_estimators: int = self.opt_params.get('n_estimators', 100)
|
|
|
|
if early_stopping_rounds is not None and dvalid is not None:
|
|
self.model = xgb.train(
|
|
params,
|
|
dtrain,
|
|
num_boost_round=n_estimators,
|
|
evals=[(dvalid, 'valid')],
|
|
early_stopping_rounds=early_stopping_rounds,
|
|
)
|
|
else:
|
|
self.model = xgb.train(
|
|
params,
|
|
dtrain,
|
|
num_boost_round=n_estimators,
|
|
)
|
|
|
|
def _train_model(
|
|
self,
|
|
x: pd.DataFrame,
|
|
y: pd.DataFrame,
|
|
x_val: pd.DataFrame | None = None,
|
|
y_val: pd.DataFrame | None = None,
|
|
) -> None:
|
|
"""
|
|
Train the underlying XGBoost model instance for the first time.
|
|
|
|
Receives already-transformed feature DataFrames from the data model.
|
|
|
|
Args:
|
|
- x (pd.DataFrame): Transformed training feature matrix.
|
|
- y (pd.DataFrame): Training target data.
|
|
- x_val (pd.DataFrame | None): Optional transformed validation features.
|
|
- y_val (pd.DataFrame | None): Optional transformed validation targets.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
self._fit_xgboost_model(x, y, x_val, y_val)
|
|
|
|
def _retrain_model(self, x: pd.DataFrame, y: pd.DataFrame | None) -> None:
|
|
"""
|
|
Retrain the XGBoost model using an internal train/validation split.
|
|
|
|
Receives an already-transformed feature DataFrame and performs a
|
|
deterministic train/validation split before delegating to `_fit_xgboost_model`.
|
|
|
|
Args:
|
|
- x (pd.DataFrame): Transformed feature matrix.
|
|
- y (pd.DataFrame | None): Target DataFrame to be used for retraining.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
if y is None:
|
|
raise ValueError('y must be provided for retrain')
|
|
|
|
# Align target with features in case the transformer dropped rows (e.g. windows/lags)
|
|
y = y.loc[x.index]
|
|
|
|
# Do train/test split internally
|
|
x_train, x_val, y_train, y_val = train_test_split(
|
|
x,
|
|
y,
|
|
test_size=0.15,
|
|
shuffle=False,
|
|
random_state=self.opt_params.get('random_state', 42),
|
|
)
|
|
|
|
if y_train is None:
|
|
raise ValueError('y_train must be provided by train_test_split')
|
|
|
|
if y_val is None:
|
|
raise ValueError('y_val must be provided by train_test_split')
|
|
|
|
self._fit_xgboost_model(x_train, y_train, x_val, y_val)
|
|
|
|
def _fit_transformer_on_full_data(self, data: pd.DataFrame) -> None:
|
|
"""
|
|
Fit the configured time series transformer on the full dataset.
|
|
|
|
The transformer receives both features and target, and its `target`
|
|
and `tags_list` attributes are populated from the wrapper before calling
|
|
`fit` so that it can correctly generate time-based features.
|
|
|
|
Args:
|
|
- data (pd.DataFrame): Full dataset including both features and target.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
if self.target is not None:
|
|
self.transformer.target = self.target
|
|
if self.model_tags is not None:
|
|
self.transformer.tags_list = self.model_tags
|
|
self.transformer.fit(data)
|
|
|
|
def _train_transformer(
|
|
self,
|
|
train_data: pd.DataFrame,
|
|
val_data: pd.DataFrame,
|
|
) -> None:
|
|
"""
|
|
Train the transformer instance using the training dataset.
|
|
|
|
The wrapper does not mutate the provided DataFrames; they are only
|
|
used to fit the transformer's internal state.
|
|
|
|
Args:
|
|
- train_data (pd.DataFrame): Full training dataset (features + target).
|
|
- val_data (pd.DataFrame): Full validation dataset (features + target),
|
|
kept for interface symmetry but not used directly in this implementation.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
self._fit_transformer_on_full_data(train_data)
|
|
|
|
def _retrain_transformer(self, data: pd.DataFrame) -> None:
|
|
"""
|
|
Retrain the transformer instance using the full dataset.
|
|
|
|
Args:
|
|
- data (pd.DataFrame): Full dataset (features + target column).
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
self._fit_transformer_on_full_data(data)
|