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