Code import - branch 8
This commit is contained in:
303
models/xgboost/xgboost_model_wrapper.py
Normal file
303
models/xgboost/xgboost_model_wrapper.py
Normal file
@@ -0,0 +1,303 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user