1352d1ac8f
- Added `evidently` to requirements for improved model evaluation. - Introduced `TrainModelResult` class with a `to_dict` method for better result handling. - Updated `train_model` method to return a comprehensive training result, including run details. - Enhanced `cleanup_run_directory` method in `DataManagerRepository` for improved resource management. - Adjusted type hints in `TrainModel` for clarity and consistency.
59 lines
2.5 KiB
Python
59 lines
2.5 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import pandas as pd
|
|
|
|
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 prepared 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.
|
|
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.
|
|
y_pred (pd.Series | None): The predicted target values for the testing dataset. Default is None.
|
|
y_train_pred (pd.Series | None): The predicted target values for the training dataset. Default is None.
|
|
mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None.
|
|
mae_val (float | None): The Mean Absolute Error (MAE) of the predictions. Default is None.
|
|
r2_val (float | None): The R-squared (R²) value of the predictions. Default is None.
|
|
equation (dict | None): The equation of the model. Default is None.
|
|
equation_path (str | None): The path to the equation file. Default is None.
|
|
run_name (str | None): The name of the MLFlow run. Default is None.
|
|
report_path (str | None): The path to the generated HTML report file. Default is None.
|
|
train_data_path (str | None): The path to the training dataset CSV file. Default is None.
|
|
test_data_path (str | None): The path to the testing dataset CSV file. Default is None.
|
|
"""
|
|
|
|
params: TrainModelParams
|
|
train_data: pd.DataFrame
|
|
val_data: pd.DataFrame
|
|
y_pred: pd.DataFrame | None = None
|
|
y_train_pred: pd.DataFrame | None = None
|
|
mse_val: float | None = None
|
|
mae_val: float | None = None
|
|
r2_val: float | None = None
|
|
equation: dict | None = None
|
|
equation_path: str | None = None
|
|
run_name: str | None = None
|
|
run_id: 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
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""
|
|
Convert TrainModelResult to a dictionary.
|
|
"""
|
|
return self.__dict__ |