Code import - branch 8

This commit is contained in:
2026-08-05 13:53:42 +00:00
commit 781dff7b9e
18 changed files with 1631 additions and 0 deletions

View File

@@ -0,0 +1,70 @@
# Linear Regression Model
Linear regression model for time series analysis, with optional polynomial features and clipping support.
## Model Parameters (`LinearRegressionModel`)
These parameters are defined in the `model.yaml` under `model.input_map` and are also represented in `schemas.yaml` under `components.schemas.model` with concrete `example` values used during automated validation:
| Parameter | Example | Description |
|--------------------|---------|-----------------------------------------------------------------------------|
| `degree` | `2` | Degree of the polynomial used to create polynomial features. |
| `interaction_only` | `true` | Whether to include only interaction features in polynomial features. |
| `verbose` | `true` | Enable verbose output during model execution. |
| `clipping_max` | `100` | Maximum value for clipping prediction output. |
| `clipping_min` | `0` | Minimum value for clipping prediction output. |
## Data Preprocessor (`DataPreprocessor`)
These parameters are defined in the `model.yaml` under `data_model.input_map` and are also represented in `schemas.yaml` under `components.schemas.data_model` with concrete `example` values used during automated validation:
| Parameter | Example | Description |
|----------------------|--------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------|
| `scaler__name` | `'Standard Scaler'` | Name of the scaler to use for scaling the data. Options: 'Standard Scaler' |
| `verbose` | `true` | Enable verbose output during model execution. |
| `steps_order` | `['Discontinuity Treatment', 'Lag Selection', 'Range Selection & Data Removal', 'Static Window Removal', 'Define Variables Limits', 'Normalization', 'Feature Creation', 'Lag Creation']` | Order of the steps to be executed in the pipeline. |
| `nan_treatment` | `'drop'` | The treatment for missing values. Options: 'drop', 'linear interpolation' |
| `lag_train` | `{'variable_name': 1}` | The lags for each variable to be applyed during training. |
| `lag_transform` | `{'variable_name': 1}` | The lags for each variable to be applyed during transformation. |
| `start_date` | `'2020-01-01 00:00:00'` | The start date for the dataset. Format: 'YYYY-MM-DD HH:MM:SS' |
| `end_date` | `'2020-01-01 00:00:00'` | The end date for the dataset. Format: 'YYYY-MM-DD HH:MM:SS' |
| `removed_intervals` | `[(2020-01-01 00:00:00, 2020-01-01 00:00:00)]` | The intervals to be removed from the dataset. |
| `static_threshold` | `10` | The number of repeated values to be considered as static. |
| `lower_limits` | `{'variable_name': 0}` | The lower limits for each variable. |
| `upper_limits` | `{'variable_name': 1}` | The upper limits for each variable. |
| `scaler_name` | `'Standard Scaler'` | The scaler name. Options: 'None', 'Standard Scaler' |
| `scaler_params` | `{'variable_name': {'mean': 0, 'variance': 1}}` | The parameters for the scaler object, if it is used. |
| `self_operations` | `['variable_name_exp_scalar']` | The operations for feature creation using the same variable. |
| `cross_operations` | `['variable_name1_mul_variable_name2']` | The operations for feature creation using two variables. |
| `created_lags` | `{'variable_name': [1, 2]}` | Variables created by lagging existing ones. |
## Lifecycle (train / retrain)
The wrapper follows the `SientiaModel` interface. Data is always passed as full DataFrames (features + target).
- **`train(train_data, val_data, target)`**
- `train_data`: full training DataFrame (features + target column).
- `val_data`: full validation DataFrame (features + target column).
- `target`: name of the target column (must exist in both).
The base class fits the transformer on the full datasets, then transforms and fits the model on transformed features and target.
- **`retrain(data)`**
- `data`: full dataset (features + target column).
Uses `self.target` set during `train()`. Call only after `train()` has been run.
Neither method mutates the passed DataFrames.
## Usage Notes
- If `variable_columns` is not set, it is inferred during `fit` as all columns except the target.
- Polynomial features are created automatically when `degree > 1`.
- `clipping` uses Q1/Q3 of the training target when replacing out-of-bounds predictions.
- The wrapper enforces that `wrapper.target` is set before training or retraining; otherwise, a `ValueError` is raised to fail fast when the configuration is incomplete.
All parameters documented above **must stay in sync** between:
- This README
- `model.yaml` (`model.input_map` and `data_model.input_map`)
- `schemas.yaml` (`components.schemas.model` and `components.schemas.data_model`)
In particular, `schemas.yaml` must contain **complete and coherent examples** for every parameter that affects runtime behaviour, since those examples are used to build `model_kwargs` and `transformer_kwargs` during validation. Parameters without examples may not be fully exercised by the automated validation flow.

View File

View File

@@ -0,0 +1,21 @@
name: linear_regression
path: linear_regression_model_wrapper
class: LinearRegressionModelWrapper
runtime: basic
requirements:
- pandas==2.3.3
- scikit-learn==1.8.0
- cloudpickle==3.1.2
- mlflow==3.8.1
- git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.11.0
- git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.8.1
version: 1.1.1
type: model
model:
class: LinearRegressionModel
external: false
path: linear_regression_model
data_model:
class: DataPreprocessor
external: true
path: sientia_model.preprocessing.preprocessing

View File

@@ -0,0 +1,265 @@
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
class LinearRegressionModel(BaseEstimator, TransformerMixin):
def __init__(
self,
target_variable: str = '',
variable_columns: list | None = None,
model_params: dict | None = None,
clipping_max: float | None = None,
clipping_min: float | None = None,
weights: dict | None = None,
degree: int = 1,
interaction_only: bool = False,
verbose: bool = False,
):
"""
Linear Regression Model for Time Series Analysis.
Args:
target_variable (str): The target variable name
variable_columns (list): The input columns names in a list
model_params (dict): The parameters used for training the model \\
clipping (dict): The lower and upper limits for the target variable to be clipped \\
*Format: {'min': min_value, 'max': max_value}*
weights (dict): The weights for the Linear Regression model \\
*Format: {'variable_name': weight}*
degree (int): The degree of the polynomial features
interaction_only (bool): Whether to include interaction terms only
verbose (bool): Whether to print verbose output
Returns:
LinearRegressionModel: The prediction model object
"""
self.target_variable: str = target_variable
self.variable_columns: list = variable_columns if variable_columns else []
self.model_params: dict = model_params if model_params else {}
self.regr = LinearRegression()
self.q1_target: float | None = None
self.q3_target: float | None = None
self.weights: dict = weights if weights else {}
self.degree: int = degree
self.interaction_only: bool = interaction_only
self.poly: PolynomialFeatures | None = None
self.verbose: bool = verbose
self.clipping: dict = {}
if clipping_max is not None:
self.clipping['max'] = clipping_max
if clipping_min is not None:
self.clipping['min'] = clipping_min
def get_regressor(self) -> LinearRegression:
"""
Get LinearRegression object
Returns:
LinearRegression objetc
"""
return self.regr
def create_poly_features(self, input_data: pd.DataFrame, fit: bool = False) -> pd.DataFrame:
"""
Function to create polynomial features
Args:
input_data (pandas.DataFrame): The data used to create the polynomial features
fit (bool): Whether to fit the polynomial features creator or just transform
Returns:
pandas.DataFrame: The data with the polynomial features
"""
original_index = input_data.index
original_shape = input_data.shape
# Create polynomial features
interaction_only = self.interaction_only
degree = self.degree
if fit:
if self.verbose:
print(f'Fitting polynomial features with degree {degree}')
self.poly = PolynomialFeatures(
degree=degree, interaction_only=interaction_only, include_bias=False
)
input_data = self.poly.fit_transform(input_data)
else:
if self.verbose:
print(f'Transforming polynomial features with degree {degree}')
if self.poly is not None:
input_data = self.poly.transform(input_data)
else:
raise ValueError(
'Polynomial transformer is not fitted. Call fit() before predict() when degree > 1.'
)
# Create new columns names
modified_columns = self.poly.get_feature_names_out(self.variable_columns)
input_data = pd.DataFrame(input_data, columns=modified_columns, index=original_index)
if self.verbose:
print(f'Data shape before transformation: {original_shape}')
print(f'Data shape after transformation: {input_data.shape}')
return input_data
def fit(self, input_data: pd.DataFrame) -> 'LinearRegressionModel':
"""
Function to fit the model
Args:
input_data (pandas.DataFrame): The data used to fit the Linear Regression model
Returns:
LinearRegressionModel: The prediction model object
"""
if self.verbose:
# Display the header
text = 'INITIATING LINEAR REGRESSION MODEL FIT'
print('\n' + '-' * len(text))
print(text)
print('-' * len(text) + '\n')
# Basic validations and inference
if not self.target_variable:
raise ValueError("'target_variable' must be set before calling fit().")
if not self.variable_columns:
# Infer all columns except target
self.variable_columns = [c for c in input_data.columns if c != self.target_variable]
# Validate required columns exist
missing_features = [c for c in self.variable_columns if c not in input_data.columns]
if missing_features:
raise ValueError(
'Training data is missing required feature columns: ' + ', '.join(missing_features)
)
if self.target_variable not in input_data.columns:
# Let pandas raise KeyError with the column name to satisfy tests too
raise KeyError(repr(self.target_variable))
X_train = input_data[self.variable_columns]
y_train = input_data[self.target_variable]
if self.verbose:
print(f'X_train shape: {X_train.shape}')
print(f'y_train shape: {y_train.shape}')
# Create polynomial features
if self.degree > 1:
X_train = self.create_poly_features(X_train, fit=True)
# Handle infinite values
X_train = X_train.replace([np.inf, -np.inf], np.nan)
# Drop columns with all null values
all_null_cols = X_train.columns[X_train.isnull().all()].tolist()
X_train = X_train.drop(columns=all_null_cols)
if self.verbose and all_null_cols:
print(f'{len(all_null_cols)} columns with all null values were dropped.')
# Get mask of rows with any NaN values
mask = X_train.notna().all(axis=1)
X_train = X_train[mask]
y_train = y_train[mask]
if self.verbose:
print(f'X_train shape after removing NaN rows: {X_train.shape}')
print(f'y_train shape after removing NaN rows: {y_train.shape}')
self.q1_target = y_train.quantile(0.25)
self.q3_target = y_train.quantile(0.75)
# Fit the model
self.regr.fit(X_train, y_train)
if self.verbose:
print('Model training completed successfully.')
# Get the weights using actual trained feature names (handles polynomial features)
# ravel() so coef_ is 1D when y was a column (e.g. DataFrame), avoiding scalar conversion errors
round_coef = np.round(self.regr.coef_, 3).ravel()
round_intercept = np.round(self.regr.intercept_, 3).ravel()
feature_names = list(getattr(self.regr, 'feature_names_in_', X_train.columns))
# Save the weights (strict=True to ensure feature_names and coefficients align)
weights = dict(zip(feature_names, [float(c) for c in round_coef], strict=True))
weights = dict(sorted(weights.items(), key=lambda item: abs(item[1]), reverse=True))
weights = {'Bias': float(round_intercept.flat[0]), **weights}
self.weights = weights
if self.verbose:
print('\nModel Weights:')
for key, value in weights.items():
print(f'{key}: {value}')
return self
def predict(self, input_data: pd.DataFrame) -> np.ndarray:
"""
Function to predict the target variable.
If clipping is True, the predictions are clipped based on the target variable quartiles.
Args:
input_data (pandas.DataFrame): The data used to predict the target variable
Returns:
numpy.ndarray: The predicted target variable
"""
if self.verbose:
# Display the header
text = 'INITIATING LINEAR REGRESSION MODEL PREDICTION'
print('\n' + '-' * len(text))
print(text)
print('-' * len(text) + '\n')
# Validate required columns exist for prediction
missing_features = [c for c in self.variable_columns if c not in input_data.columns]
if missing_features:
raise ValueError(
'Prediction data is missing required feature columns: '
+ ', '.join(missing_features)
)
X_test = input_data[self.variable_columns]
if self.verbose:
print(f'X_test shape: {X_test.shape}')
# Create polynomial features
if self.degree > 1:
X_test = self.create_poly_features(X_test, fit=False)
X_test = X_test.replace([np.inf, -np.inf], np.nan)
model_features = self.regr.feature_names_in_
# Ensure required model features exist
missing_model_features = [c for c in model_features if c not in X_test.columns]
if missing_model_features:
raise ValueError(
'Prediction data is missing required transformed feature columns: '
+ ', '.join(missing_model_features)
)
# Get mask of rows with any NaN values
mask = X_test.notna().all(axis=1)
X_test = X_test[mask]
if self.verbose:
print(f'X_test shape after removing NaN rows: {X_test.shape}')
# For simplicity and to keep behavior consistent with prior tests, drop NaN rows before prediction
y_pred = self.regr.predict(X_test[model_features])
if self.clipping:
for i in range(len(y_pred)):
if y_pred[i] > self.clipping['max']:
y_pred[i] = self.q3_target
elif y_pred[i] < self.clipping['min']:
y_pred[i] = self.q1_target
return y_pred

View File

@@ -0,0 +1,252 @@
from typing import Any
import pandas as pd
from mlflow.models.signature import ModelSignature
from sientia_model.wrappers.sientia_model import SientiaModel
class LinearRegressionModelWrapper(SientiaModel):
"""
Wrapper around `LinearRegressionModel` that conforms to the SientiaModel lifecycle.
The wrapper coordinates the external `DataPreprocessor` and the underlying
linear regression model so that both training and inference respect the
Sientia runtime and storage conventions.
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 by the base
SientiaModel; the linear regression wrapper itself reads target and
feature definitions from `model.yaml` and the data.
Return:
None
"""
def _load_model(self, **kwargs: Any) -> None:
"""
Apply post-construction configuration to the underlying model instance.
This hook is currently a no-op but is kept for compatibility with the
SientiaModel interface and future extensions.
Args:
- **kwargs (Any): Keyword arguments that were passed to the model
constructor (`model_kwargs` from `load_stack`).
Return:
None
"""
pass
def _load_transformer(self, **kwargs: Any) -> None:
"""
Apply post-construction configuration to the underlying transformer.
This hook is currently a no-op but is kept for compatibility with the
SientiaModel interface and future extensions.
Args:
- **kwargs (Any): Keyword arguments that were passed to the transformer
constructor (`transformer_kwargs` from `load_stack`).
Return:
None
"""
if self.transformer.removed_intervals is None:
return
parsed_intervals = []
for interval in self.transformer.removed_intervals:
# intervals came in format "(start_date, end_date)"
# we need to convert to a tuple of strings
parsed_intervals.append((interval[0], interval[1]))
self.transformer.removed_intervals = parsed_intervals
def _load_context(self, context: Any) -> None:
"""
Optionally load runtime context information from MLflow.
The linear regression wrapper does not need any context information,
so this method intentionally performs no work.
Args:
- context (Any): MLflow context object provided during model loading.
Return:
None
"""
pass
def _predict(self, data: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, Any]]:
"""
Generate predictions using the underlying `LinearRegressionModel`.
Args:
- data (pd.DataFrame): Input features for prediction.
Return:
tuple[pd.DataFrame, dict[str, Any]]: DataFrame containing one prediction
column named after the wrapper target, plus an empty metadata dictionary.
"""
array_predictions = self.model.predict(data)
predictions_df = pd.DataFrame(array_predictions, columns=[self.target], index=data.index)
return predictions_df, {}
def _transform(self, data: pd.DataFrame) -> tuple[pd.DataFrame, dict[str, Any]]:
"""
Execute transformation for the given input data.
When a transformer is configured, its `predict` method is used to
transform the incoming features; this is how the external
`DataPreprocessor` is applied before calling the model.
Args:
- data (pd.DataFrame): Input features as a pandas DataFrame.
Return:
tuple[pd.DataFrame, dict[str, Any]]: Transformed DataFrame and an
empty metadata dictionary.
"""
transformed_data = self.transformer.predict(data)
return transformed_data, {}
def _fit_linear_model(self, x: pd.DataFrame, y: pd.DataFrame) -> None:
"""
Fit the underlying `LinearRegressionModel` with the provided training data.
Args:
- x (pd.DataFrame): Training feature matrix.
- y (pd.DataFrame): Training target data as a one-column DataFrame.
Return:
None
"""
if self.target is None:
raise ValueError('wrapper.target must be set when providing target data for training.')
target = self.target
# Use preprocessor output as-is when it already includes the target; otherwise concat x and y
if y is not None and target not in x.columns:
# Align target with features before potential index changes
y = y.loc[x.index]
data = pd.concat([x, y], axis=1)
else:
data = x.copy()
# Set target variable name on the model
self.model.target_variable = target
# Fit the model with the training dataset
self.model.fit(data)
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 model instance using transformed feature data.
Args:
- x (pd.DataFrame): Transformed training feature matrix.
- y (pd.DataFrame): Training target DataFrame.
- x_val (pd.DataFrame | None): Transformed validation feature matrix.
- y_val (pd.DataFrame | None): Validation target DataFrame.
Return:
None
"""
self._fit_linear_model(x, y)
def _retrain_model(self, x: pd.DataFrame, y: pd.DataFrame | None) -> None:
"""
Retrain the underlying model instance with new transformed data.
Args:
- x (pd.DataFrame): Transformed feature matrix as pandas DataFrame.
- y (pd.DataFrame | None): Target DataFrame; must be provided.
Return:
None
"""
if y is None:
raise ValueError('y must be provided for retrain')
self._fit_linear_model(x, y)
def _fit_transformer_on_full_data(self, data: pd.DataFrame) -> None:
"""
Fit the external transformer with the full dataset (features + target).
The wrapper sets `target_variable` and `input_variables` on the
transformer so it can correctly access the target column and the
feature list (for example, the `DataPreprocessor`).
Args:
- data (pd.DataFrame): Full dataset (features + target column).
Return:
None
"""
if self.transformer is not None:
self.transformer.target_variable = self.target
self.transformer.input_variables = self.model_tags
self.transformer.fit(data)
def _train_transformer(
self,
train_data: pd.DataFrame,
val_data: pd.DataFrame,
) -> None:
"""
Train the external transformer using the training data.
Args:
- train_data (pd.DataFrame): Full training dataset (features + target).
- val_data (pd.DataFrame): Full validation dataset (features + target),
included for signature compatibility (not used directly).
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)
def _store_model(
self,
final_name: str,
signature: ModelSignature,
pip_requirements: list[str],
code_path: list[str],
) -> None:
"""
Store the model artifacts in MLflow without adding extra payloads.
This wrapper does not need to persist custom artifacts beyond those
handled by the base SientiaModel, so this implementation is a no-op.
Args:
- final_name (str): Final model name used in MLflow and Model Registry.
- signature (ModelSignature): MLflow model signature defining input and output schema.
- pip_requirements (list[str]): Python package requirements for the model.
- code_path (list[str]): List of code paths packaged with the model.
Return:
None
"""
pass

View File

@@ -0,0 +1,206 @@
openapi: 3.1.0
info:
title: Linear Regression model API
version: 1.0.0
components:
schemas:
model:
type: object
properties:
degree:
type: integer
default: 1
example: 2
description: "Degree of the polynomial used to create polynomial features."
interaction_only:
type: boolean
default: false
example: true
description: "Whether to include only interaction features in polynomial features."
verbose:
type: boolean
default: false
example: true
description: "Enable verbose output during model execution."
clipping_max:
type: integer
default: null
example: 100
description: "Maximum value for clipping prediction output."
clipping_min:
type: integer
default: null
example: 0
description: "Minimum value for clipping prediction output."
data_model:
type: object
properties:
verbose:
type: boolean
default: false
example: true
description: "Enable verbose output during model execution."
StepsList: &steps_list
- "Discontinuity Treatment"
- "Lag Selection"
- "Range Selection & Data Removal"
- "Static Window Removal"
- "Define Variables Limits"
- "Normalization"
- "Feature Creation"
- "Lag Creation"
steps_order:
type: array
items:
type: string
enum: *steps_list
required: *steps_list
default: *steps_list
example: *steps_list
uniqueItems: true
description: "Order of the steps to be executed in the pipeline."
nan_treatment:
type: string
enum: ["drop", "linear interpolation"]
default: "drop"
example: "drop"
description: "The treatment for missing values."
lag_train:
type: object
x-feature: true
additionalProperties:
type: integer
example: { "col_2": 1 }
default: {}
description: "The lags for each variable to be applyed during training. Format: {'variable_name': lag}"
lag_transform:
type: object
x-feature: true
additionalProperties:
type: integer
example: { "col_2": 1 }
default: {}
description: "The lags for each variable to be applyed during transformation. Format: {'variable_name': lag}"
start_date:
type: string
format: "date-time"
default: null
example: "2020-01-01 00:00:00"
description: "The start date for the dataset. Format: 'YYYY-MM-DD HH:MM:SS'"
end_date:
type: string
format: "date-time"
default: null
example: "2020-04-25 00:00:00"
description: "The end date for the dataset. Format: 'YYYY-MM-DD HH:MM:SS'"
removed_intervals:
type: array
items:
type: array
items:
type: string
format: "date-time"
example:
- ["2020-01-01 00:00:00", "2020-01-02 00:00:00"]
default: []
description: "The intervals to be removed from the dataset. Format: [[start_date, end_date], ...]"
static_threshold:
type: integer
default: null
example: 10
description: "The number of repeated values to be considered as static"
lower_limits:
type: object
x-feature: true
additionalProperties:
type: number
example:
{ "col_2": 0 }
default: {}
description: "The lower limits for each variable. Format: {'variable_name': limit}"
upper_limits:
type: object
x-feature: true
additionalProperties:
type: number
example:
{ "col_2": 100 }
default: {}
description: "The upper limits for each variable. Format: {'variable_name': limit}"
scaler_name:
type: string
enum:
- "Standard Scaler"
default: "Standard Scaler"
example: "Standard Scaler"
description: "The scaler name."
scaler_params:
type: object
additionalProperties:
type: object
properties:
mean:
type: number
variance:
type: number
required:
- mean
- variance
example:
{ "col_2": { "mean": 0, "variance": 1 }}
default: {}
description: "The parameters for the scaler object."
self_operations:
type: array
items:
type: object
properties:
variable:
type: string
x-feature: true
operation:
type: string
enum: ["exp", "pow", "log", "root"]
scalar:
type: number
required:
- variable
- operation
- scalar
example: [{ "variable": "col_2", "operation": "pow", "scalar": 2 }]
default: []
description: "The operations for feature creation using the same variable."
cross_operations:
type: array
items:
type: object
properties:
variable1:
type: string
x-feature: true
variable2:
type: string
x-feature: true
operation:
type: string
enum: ["*", "/"]
required:
- variable1
- variable2
- operation
example:
[{ "variable1": "col_2", "variable2": "col_3", "operation": "*" }]
default: []
description: "The operations for feature creation using two variables."
created_lags:
type: object
x-feature: true
additionalProperties:
type: array
items:
type: integer
example:
{ "col_2": [1, 2] }
default: {}
description: "Variables created by lagging existing ones."

72
models/xgboost/README.md Normal file
View File

@@ -0,0 +1,72 @@
# XGBoost Model
XGBoost model for time series regression/classification, with a built-in `TimeSeriesPreprocessor` for rolling features and scaling.
## Model Parameters (`Booster`)
These parameters are defined in the `model.yaml` under `opt_params_map` (which configure the XGBoost booster training) and are also represented in `schemas.yaml` under `components.schemas.opt_params` with concrete `example` values used during automated validation:
| Parameter | Example | Description |
|-------------------------|----------|-----------------------------------------------------------------------------|
| `tree_method` | `"hist"` | Tree construction algorithm. Options: `'hist'`, `'exact'`, `'approx'`, etc. |
| `device` | `"cuda"` | Device for computation (`cuda` or `cpu`). |
| `learning_rate` | `0.3` | Step size shrinkage to prevent overfitting. |
| `n_estimators` | `100` | Number of boosting rounds. |
| `max_depth` | `32` | Maximum depth of a tree. |
| `subsample` | `0.8` | Subsample ratio of the training instances (01). |
| `colsample_bytree` | `0.8` | Subsample ratio of columns when constructing each tree (01). |
| `min_child_weight` | `5` | Minimum sum of instance weight in a child node. |
| `random_state` | `42` | Seed for XGBoost random number generation. Set this to make training fully reproducible. |
| `early_stopping_rounds` | `None` | If set, enables early stopping after N rounds with no improvement on validation. Requires an internal 15% validation split. |
## Data Model Parameters (`TimeSeriesPreprocessor`)
These parameters are defined in the `model.yaml` under `data_model.input_map` and are also represented in `schemas.yaml` under `components.schemas.data_model` with concrete `example` values used during automated validation:
| Parameter | Example | Description |
|-------------------|------------------------------------------------------|-----------------------------------------------------------------------------|
| `scaler_method` | `"MinMax"` | Scaling strategy, either `'MinMax'`, `'Standard'` or None for no scaling. |
| `window_size` | `3` | Rolling window size used to generate lagged statistics for the target. Zero disables rolling features. |
| `use_filtering` | `true` | Whether to filter rows based on a CI column during `fit` when `arpr_config` is provided. |
| `transform_mode` | `"all"` | Prediction mode, `'all'` to keep all rows or `'latest'` to return only the last row. |
| `arpr_config` | `{"ci_col": "ci", "ar_col": "ar", "pr_col": "pr"}` | Optional configuration mapping `'ci_col'`, `'ar_col'` and `'pr_col'` names used for AR/PR substitution. |
### Rolling Features
When `window_size > 0`, the preprocessor generates rolling features for the target: `rolling_mean`, `rolling_max`, `rolling_min`, `rolling_std`.
### AR/PR Substitution
If `arpr_config` contains `ar_col`, `pr_col`, and `ci_col`, rows where `ci_col == 0` have `ar_col` replaced with `pr_col`.
## Lifecycle (train / retrain)
The wrapper follows the `SientiaModel` interface. Data is always passed as full DataFrames (features + target).
- **`train(train_data, val_data, target)`**
- `train_data`: full training DataFrame (features + target column).
- `val_data`: full validation DataFrame (features + target column).
- `target`: name of the target column (must exist in both).
The base class fits the transformer on the full datasets, then transforms and fits the model on transformed features and target.
- **`retrain(data)`**
- `data`: full dataset (features + target column).
Uses `self.target` set during `train()`. Call only after `train()` has been run. Retrain uses an internal 15% validation split when `early_stopping_rounds` is set.
Neither method mutates the passed DataFrames.
## Data Format
- Timestamp must be represented only on the DataFrame index (do not provide `timestamp` or `Timestamp` as regular columns).
- Input and output frames used in `transform()`/`predict()` must use `DatetimeIndex` with datetime64 dtype.
- Runtime validation checks that output max timestamp stays aligned with the corresponding input max timestamp for both `transform` and `predict` flows.
- Training uses an internal 15% validation split when `early_stopping_rounds` is set.
- **Data Alignment**: The wrapper automatically aligns the target variable with the features after transformation. This ensures that if the preprocessor drops rows (e.g., due to rolling windows or lags), the model training remains consistent.
All fields documented above should remain consistent across:
- This README
- `model.yaml` (`opt_params_map` and `data_model.input_map`)
- `schemas.yaml` (`components.schemas.opt_params` and `components.schemas.data_model`)
Examples defined in `schemas.yaml` are used directly during runtime validation to build `opt_params` and `transformer_kwargs`. Parameters without examples may not be fully exercised in the automated validation flow, so authors should always provide realistic and complete examples for all relevant fields.

View File

24
models/xgboost/index.yaml Normal file
View File

@@ -0,0 +1,24 @@
name: xgboost
path: xgboost_model_wrapper
class: XGBoostWrapper
runtime: xgboost
requirements:
- mlflow==3.8.1
- pathspec==1.0.4
- pandas==2.3.3
- xgboost ==3.1.2
- cloudpickle==3.1.2
- rich==14.3.2
- joblib==1.5.3
- git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.11.0
- git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.8.1
version: 1.1.1
type: model
model:
external: true
class: Booster
path: xgboost
data_model:
class: TimeSeriesPreprocessor
external: false
path: xgboost_data_model

107
models/xgboost/schemas.yaml Normal file
View File

@@ -0,0 +1,107 @@
openapi: 3.1.0
info:
title: XGBoost model API
version: 1.0.0
components:
schemas:
data_model:
type: object
properties:
scaler_method:
type: string
enum: ["MinMax", "Standard Scaler"]
default: null
example: "MinMax"
description: "Scaling strategy."
window_size:
type: integer
default: 0
example: 3
description: "Rolling window size used to generate lagged statistics for the target."
use_filtering:
type: boolean
default: false
example: false
description: "Whether to filter rows based on a CI column during `fit` (expects a binary CI column)."
transform_mode:
type: string
enum: ["all", "latest"]
default: "all"
example: "all"
description: "Prediction mode, `'all'` to keep all rows or `'latest'` to return only the last row."
arpr_config:
type: object
properties:
ci_col:
type: string
x-feature: true
ar_col:
type: string
x-feature: true
pr_col:
type: string
x-feature: true
example:
ci_col: "col_1"
ar_col: "col_2"
pr_col: "col_3"
description: "Optional configuration mapping for AR/PR substitution."
opt_params:
type: object
properties:
tree_method:
type: string
enum: ["hist", "exact", "approx"]
default: "hist"
example: "hist"
description: "Tree construction algorithm."
device:
type: string
enum: ["cuda", "cpu"]
default: "cuda"
example: "cuda"
description: "Device for computation."
learning_rate:
type: number
default: 0.3
example: 0.3
description: "Step size shrinkage to prevent overfitting."
n_estimators:
type: integer
default: 100
example: 100
description: "Number of boosting rounds."
max_depth:
type: integer
default: 32
example: 32
description: "Maximum depth of a tree."
subsample:
type: number
minimum: 0.0
maximum: 1.0
default: 0.8
example: 0.8
description: "Subsample ratio of the training instances."
colsample_bytree:
type: number
minimum: 0.0
maximum: 1.0
default: 0.8
example: 0.8
description: "Subsample ratio of columns when constructing each tree."
min_child_weight:
type: integer
default: 5
example: 5
description: "Minimum sum of instance weight in a child node."
random_state:
type: integer
default: 42
example: 42
description: "Seed for XGBoost random number generation."
early_stopping_rounds:
type: integer
default: null
example: 100
description: "If set, enables early stopping."

View File

@@ -0,0 +1,234 @@
"""
Standardized preprocessor for time series models used with XGBoost wrappers.
Args:
- None
Return:
None
"""
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import MinMaxScaler, StandardScaler
class TimeSeriesPreprocessor(BaseEstimator, TransformerMixin):
def __init__(
self,
target: str | None = None,
tags_list: list[str] | None = None,
scaler_method: str | None = None,
window_size: int = 0,
use_filtering: bool = False,
transform_mode: str = 'all',
arpr_config: dict[str, str] | None = None,
) -> None:
"""
Initialize the time series preprocessor for feature engineering and scaling.
Args:
- target (str | None): Name of the target column. If None, it must be
set before calling `fit`.
- tags_list (list[str] | None): Explicit list of feature columns. If
None, inferred during `fit` by excluding target (and CI column when
AR/PR filtering is enabled).
- scaler_method (str | None): Scaling strategy, either `'MinMax'`,
`'Standard'` or None for no scaling.
- window_size (int): Rolling window size used to generate lagged
statistics for the target. Zero disables rolling features.
- use_filtering (bool): Whether to filter rows based on a CI column
during `fit` when `arpr_config` is provided.
- transform_mode (str): Prediction mode, `'all'` to keep all rows or
`'latest'` to return only the last row.
- arpr_config (dict[str, str] | None): Optional configuration mapping
`'ci_col'`, `'ar_col'` and `'pr_col'` names used for AR/PR substitution.
Return:
None
"""
self.target = target
self.tags_list = tags_list
self.scaler_method = scaler_method
self.window_size = window_size
self.use_filtering = use_filtering
self.transform_mode = transform_mode
if arpr_config is None:
arpr_config = {}
self.ci_col = arpr_config.get('ci_col')
self.ar_col = arpr_config.get('ar_col')
self.pr_col = arpr_config.get('pr_col')
if self.ci_col is None or self.ar_col is None or self.pr_col is None:
self.ar_filter = False
else:
self.ar_filter = True
# Internal state
self.scaler = None
def _adjust_time_index(self, x: pd.DataFrame) -> pd.DataFrame:
"""
Standardize the time index so it is compatible with XGBoost expectations.
Args:
- x (pd.DataFrame): Input DataFrame that may contain a time column.
Return:
pd.DataFrame: Copy of the input with the index set to `Timestamp` or
`timestamp` when present; otherwise the original index is preserved.
"""
df = x.copy()
time_cols = ['Timestamp', 'timestamp']
for col in time_cols:
if col in df.columns:
df[col] = pd.to_datetime(df[col])
df = df.set_index(col)
break
return df
def _generate_rolling_features(self, df: pd.DataFrame) -> pd.DataFrame:
"""
Generate rolling statistics for the target column.
The method computes rolling mean, max, min and standard deviation for
the configured target when `window_size` is greater than zero.
Args:
- df (pd.DataFrame): Input DataFrame indexed by time and containing
the target column.
Return:
pd.DataFrame: DataFrame with new rolling feature columns added. When
the window is larger than the number of rows, the result may be empty
after `dropna`.
"""
if self.target is None:
raise ValueError('target is not set. Set it before calling fit.')
if self.window_size > 0:
df = df.sort_index()
if df.shape[0] < self.window_size:
# If we don't have enough data for the window, we might have issues
# but we'll try to calculate what we can.
pass
df['rolling_mean'] = df[self.target].rolling(window=self.window_size).mean()
df['rolling_max'] = df[self.target].rolling(window=self.window_size).max()
df['rolling_min'] = df[self.target].rolling(window=self.window_size).min()
df['rolling_std'] = df[self.target].rolling(window=self.window_size).std()
df = df.dropna()
return df
def _apply_ar_substitution(self, df: pd.DataFrame) -> pd.DataFrame:
"""
Apply AR/PR substitution based on the configured CI (confidence) column.
When AR/PR filtering is enabled and the CI column is present, rows
where `ci_col == 0` have the AR column replaced by the PR column.
Args:
- df (pd.DataFrame): Input DataFrame containing AR, PR and CI columns.
Return:
pd.DataFrame: DataFrame with AR values updated where CI indicates PR
substitution, or unchanged when filtering is disabled or misconfigured.
"""
if self.ar_filter:
ar_col = self.ar_col
pr_col = self.pr_col
ci_col = self.ci_col
if all([ar_col, pr_col, ci_col]) and ci_col in df.columns:
mask = df[ci_col] == 0
if mask.any():
df.loc[mask, ar_col] = df.loc[mask, pr_col]
return df
def fit(self, x: pd.DataFrame) -> None:
"""
Fit the preprocessor on the full dataset.
This step may filter rows, generate rolling features, infer `tags_list`
when it is not provided, and fit the underlying scaler if configured.
Args:
- x (pd.DataFrame): Full dataset containing the target and feature
columns, and optionally CI/AR/PR columns.
Return:
None
"""
if self.target is None:
raise ValueError('target is not set. Set it before calling fit.')
df = self._adjust_time_index(x)
# Optional filtering during fit
if self.use_filtering and self.ar_filter:
ci_col = self.ci_col
if ci_col in df.columns:
df = df[df[ci_col] == 1]
# Generate rolling features if needed
df = self._generate_rolling_features(df)
# Infer tags_list if not provided
if self.tags_list is None:
exclude_cols = [self.target]
if self.ar_filter and self.ci_col is not None:
exclude_cols.append(self.ci_col)
self.tags_list = [c for c in df.columns if c not in exclude_cols]
# Fit scaler
if self.scaler_method == 'MinMax':
self.scaler = MinMaxScaler()
elif self.scaler_method == 'Standard':
self.scaler = StandardScaler()
if self.scaler:
# Only fit on tags_list_
self.scaler.fit(df[self.tags_list])
def transform(self, x: pd.DataFrame) -> pd.DataFrame:
"""
Transform input data using the fitted preprocessor.
The transformation applies the same steps used during `fit` (time
indexing, AR/PR substitution and rolling features), then selects the
feature columns and optionally applies scaling and latest-row selection.
Args:
- x (pd.DataFrame): Input data to transform.
Return:
pd.DataFrame: Transformed feature matrix ready to be passed to the
model wrapper.
"""
if self.tags_list is None:
raise ValueError('tags_list is not set. Call fit first.')
df = self._adjust_time_index(x)
# Apply AR/PR substitution if configured
df = self._apply_ar_substitution(df)
# Generate rolling features
df = self._generate_rolling_features(df)
# Select features
df = df[self.tags_list]
# Apply scaling
if self.scaler:
data_scaled = self.scaler.transform(df)
df = pd.DataFrame(data=data_scaled, columns=self.tags_list, index=df.index)
# Transform mode: only latest sample
if self.transform_mode == 'latest':
df = df.tail(1)
return df

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