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