253 lines
8.4 KiB
Python
253 lines
8.4 KiB
Python
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
|