Code import - branch 8

models/linear_regression/linear_regression_model_wrapper.py

@@ -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