SIENTIAPDE-1249: Implement data models for Model Manager and add unit tests. This commit introduces data transfer objects (DTOs) and model classes for experiment status, training parameters, and training results, along with corresponding unit tests to ensure their correct behavior.
This commit is contained in:
55
model_manager/utils/models/train_model_result.py
Normal file
55
model_manager/utils/models/train_model_result.py
Normal file
@@ -0,0 +1,55 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pandas as pd
|
||||
from sientia.linear_models import LinearRegressionModel
|
||||
from sientia.preprocessing import DataPreprocessor
|
||||
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainModelResult:
|
||||
"""
|
||||
A data container for storing the results of a machine learning training process.
|
||||
|
||||
This dataclass encapsulates all outputs from the training pipeline, including
|
||||
the trained model, datasets, evaluation metrics, and paths to generated artifacts.
|
||||
It is used to pass results between activities in the training workflow.
|
||||
|
||||
Attributes:
|
||||
params (TrainModelParams): The parameters used to train the model.
|
||||
process_data (DataPreprocessor): The data preprocessor object used to process the input data.
|
||||
X_train (pd.DataFrame): The training dataset features.
|
||||
X_test (pd.DataFrame): The testing dataset features.
|
||||
y_train (pd.DataFrame): The training dataset target values.
|
||||
y_test (pd.DataFrame): The testing dataset target values.
|
||||
regr (LinearRegressionModel): The trained linear regression model.
|
||||
scaler_dict (dict): A dictionary containing the scalers used to scale the features and target values.
|
||||
y_pred (Optional[pd.Series]): The predicted target values for the testing dataset. Default is None.
|
||||
mse_val (Optional[float]): The Mean Squared Error (MSE) of the predictions. Default is None.
|
||||
mae_val (Optional[float]): The Mean Absolute Error (MAE) of the predictions. Default is None.
|
||||
r2_val (Optional[float]): The R-squared (R²) value of the predictions. Default is None.
|
||||
run_name (Optional[str]): The name of the MLFlow run. Default is None.
|
||||
report_path (Optional[str]): The path to the generated HTML report file. Default is None.
|
||||
train_data_path (Optional[str]): The path to the training dataset CSV file. Default is None.
|
||||
test_data_path (Optional[str]): The path to the testing dataset CSV file. Default is None.
|
||||
run_dir (Optional[str]): The path to the run directory containing all artifacts. Default is None.
|
||||
"""
|
||||
|
||||
params: TrainModelParams
|
||||
process_data: DataPreprocessor
|
||||
X_train: pd.DataFrame
|
||||
X_test: pd.DataFrame
|
||||
y_train: pd.DataFrame
|
||||
y_test: pd.DataFrame
|
||||
regr: LinearRegressionModel
|
||||
scaler_dict: dict
|
||||
y_pred: pd.Series | None = None
|
||||
mse_val: float | None = None
|
||||
mae_val: float | None = None
|
||||
r2_val: float | None = None
|
||||
run_name: str | None = None
|
||||
report_path: str | None = None
|
||||
train_data_path: str | None = None
|
||||
test_data_path: str | None = None
|
||||
run_dir: str | None = None
|
||||
Reference in New Issue
Block a user