257 lines
10 KiB
Python
257 lines
10 KiB
Python
"""
|
|
Training repository for ML model training operations.
|
|
|
|
This module provides the core training logic for machine learning models,
|
|
including data preprocessing, model training, and post-training calculations.
|
|
Migrated from laborious/utils/train_model_utils.py.
|
|
"""
|
|
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.operations.df_preprocessor import load_data
|
|
from sientia_do.operations.normalization import MinMaxScaler, Z_Scaler
|
|
|
|
from model_manager.sientia.metrics import mae, mse, r2
|
|
from model_manager.sientia.models import DataPreprocessor, LinearRegressionModel
|
|
from model_manager.sientia.utils import split_train_test
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
|
|
class TrainingRepository:
|
|
"""
|
|
Repository for machine learning model training operations.
|
|
|
|
This class encapsulates the core logic for training ML models, migrated from
|
|
laborious/utils/train_model_utils.py. Follows the same pattern as MLFlowRepository
|
|
with instance methods and logger integration.
|
|
|
|
Attributes:
|
|
logger (Logger): Logger instance for observability and debugging
|
|
"""
|
|
|
|
def __init__(self, logger: Logger):
|
|
"""
|
|
Initialize TrainingRepository with logger.
|
|
|
|
Args:
|
|
logger: Logger instance for observability
|
|
"""
|
|
self.logger = logger
|
|
|
|
def train(self, uploaded_file: BytesIO, params: TrainModelParams) -> TrainModelResult:
|
|
"""
|
|
Train a machine learning model using the provided file and parameters.
|
|
|
|
This method orchestrates the training pipeline:
|
|
1. Load data from BytesIO file
|
|
2. Initialize and fit data preprocessor
|
|
3. Transform data and validate
|
|
4. Split into train/test sets
|
|
5. Initialize scaler dictionary
|
|
6. Train LinearRegression model
|
|
|
|
Args:
|
|
uploaded_file: BytesIO object containing training data (CSV format)
|
|
params: Training parameters (TrainModelParams)
|
|
|
|
Returns:
|
|
TrainModelResult: Object containing trained model, processed data,
|
|
train/test splits, and scaler dictionary
|
|
|
|
Raises:
|
|
ValueError: If transformed data is empty
|
|
Exception: If data loading, preprocessing, or training fails
|
|
"""
|
|
self.logger.info('Loading data from BytesIO file')
|
|
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
|
|
self.logger.info('Initializing and fitting data preprocessor')
|
|
process_data = self._init_data_preprocessor(params)
|
|
process_data.fit(data)
|
|
data_view = process_data.transform(data)
|
|
|
|
if len(data_view) <= 0:
|
|
raise ValueError('Data view is empty after transformation')
|
|
|
|
self.logger.info('Splitting data into train/test sets')
|
|
x_train, x_test, y_train, y_test = split_train_test(
|
|
data_view[params.variable_columns],
|
|
data_view[params.target_variable],
|
|
train_size=params.train_size / 100,
|
|
shuffle=params.shuffle,
|
|
random_state=42,
|
|
)
|
|
|
|
self.logger.info('Preparing training data')
|
|
data_train = pd.concat([x_train, y_train], axis=1)
|
|
scaler_dict = self._init_scaler_dict(process_data, params)
|
|
|
|
self.logger.info('Training linear regression model')
|
|
regr = LinearRegressionModel(
|
|
target_variable=params.target_variable,
|
|
variable_columns=params.variable_columns,
|
|
)
|
|
|
|
regr.fit(data_train)
|
|
|
|
return TrainModelResult(
|
|
params=params,
|
|
process_data=process_data,
|
|
x_train=x_train,
|
|
x_test=x_test,
|
|
y_train=y_train,
|
|
y_test=y_test,
|
|
regr=regr,
|
|
scaler_dict=scaler_dict,
|
|
)
|
|
|
|
def after_train_calculation(
|
|
self, params: TrainModelParams, tmr: TrainModelResult
|
|
) -> TrainModelResult:
|
|
"""
|
|
Perform post-training calculations: predictions, denormalization, and metrics.
|
|
|
|
This method completes the training pipeline by:
|
|
1. Making predictions on test set
|
|
2. Denormalizing all data (if scaler was used)
|
|
3. Reordering data by index
|
|
4. Calculating evaluation metrics (MSE, MAE, R²)
|
|
|
|
Args:
|
|
params: Training parameters used during model training
|
|
tmr: Result object from training
|
|
|
|
Returns:
|
|
TrainModelResult: Updated result with predictions, denormalized data,
|
|
and metrics (mse_val, mae_val, r2_val)
|
|
"""
|
|
self.logger.info('Making predictions on test set')
|
|
y_pred_array = tmr.regr.predict(tmr.x_test)
|
|
|
|
if params.use_scaler:
|
|
scaler = tmr.process_data.get_scaler()
|
|
self.logger.info('Denormalizing features')
|
|
|
|
# If using custom scaler with denormalize_* helpers
|
|
if hasattr(scaler, 'denormalize_single_input'):
|
|
for col in params.variable_columns:
|
|
tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
|
|
tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
|
|
|
|
self.logger.info('Denormalizing target variable')
|
|
tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
|
|
tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
|
|
y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
|
|
else:
|
|
# Fallback for sklearn StandardScaler: only inverse-transform features
|
|
feature_cols = getattr(
|
|
tmr.process_data, 'feature_names_order', params.variable_columns
|
|
)
|
|
# Ensure columns are in the same order used during fit
|
|
x_train_features = tmr.x_train[feature_cols]
|
|
x_test_features = tmr.x_test[feature_cols]
|
|
|
|
tmr.x_train[feature_cols] = scaler.inverse_transform(x_train_features)
|
|
tmr.x_test[feature_cols] = scaler.inverse_transform(x_test_features)
|
|
# Target was not scaled with StandardScaler in preprocessing; leave y as-is
|
|
|
|
self.logger.info('Adding index to predictions')
|
|
tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
|
|
tmr.y_pred.name = f'{params.target_variable}_pred'
|
|
|
|
self.logger.info('Reordering all data by index')
|
|
tmr.x_train = tmr.x_train.sort_index()
|
|
tmr.x_test = tmr.x_test.sort_index()
|
|
tmr.y_train = tmr.y_train.sort_index()
|
|
tmr.y_test = tmr.y_test.sort_index()
|
|
tmr.y_pred = tmr.y_pred.sort_index()
|
|
|
|
self.logger.info('Calculating evaluation metrics')
|
|
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
|
|
|
tmr.mse_val = round(
|
|
mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
|
2,
|
|
)
|
|
|
|
tmr.mae_val = round(
|
|
mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
|
2,
|
|
)
|
|
|
|
tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
|
|
return tmr
|
|
|
|
def _init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
|
|
"""
|
|
Initialize dictionary containing scaling parameters for features and target.
|
|
|
|
This method extracts scaling parameters from the fitted scaler to enable
|
|
denormalization of predictions and debugging of the normalization process.
|
|
|
|
Args:
|
|
process_data: Fitted DataPreprocessor object with scaler
|
|
params: Training parameters including scaler configuration
|
|
|
|
Returns:
|
|
dict: Scaling parameters for each feature and target variable.
|
|
Structure depends on scaler type:
|
|
- MinMaxScaler: {'feature': {'min': float, 'max': float}, ...}
|
|
- Z_Scaler: Dictionary from scaler.create_dict()
|
|
- Empty dict: If no scaler is used
|
|
|
|
Raises:
|
|
AttributeError: If scaler doesn't have expected attributes
|
|
"""
|
|
scaler_dict = {}
|
|
|
|
if params.use_scaler:
|
|
scaler = process_data.get_scaler()
|
|
|
|
if isinstance(scaler, MinMaxScaler):
|
|
# Extract min/max for each feature
|
|
for i, col in enumerate(params.variable_columns):
|
|
scaler_dict[col] = {'min': scaler.x_min[i], 'max': scaler.x_max[i]}
|
|
|
|
# Extract min/max for target variable
|
|
scaler_dict[params.target_variable] = {
|
|
'min': scaler.y_min,
|
|
'max': scaler.y_max,
|
|
}
|
|
|
|
elif isinstance(scaler, Z_Scaler):
|
|
scaler_dict = scaler.create_dict()
|
|
|
|
return scaler_dict
|
|
|
|
def _init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
|
|
"""
|
|
Initialize DataPreprocessor with training parameters.
|
|
|
|
Args:
|
|
params: Training parameters containing preprocessor configuration
|
|
|
|
Returns:
|
|
DataPreprocessor: Configured preprocessor ready for fitting
|
|
"""
|
|
# Create lag dictionaries for each variable
|
|
lag_train_dict = dict.fromkeys(params.variable_columns, params.lag_train)
|
|
lag_val_dict = dict.fromkeys(params.variable_columns, params.lag_val)
|
|
|
|
return DataPreprocessor(
|
|
target_variable=params.target_variable,
|
|
input_columns=params.variable_columns,
|
|
lag_train=lag_train_dict,
|
|
lag_transform=lag_val_dict,
|
|
static_threshold=1 if params.rem_static_win else None,
|
|
low_lim=params.low_lim,
|
|
upp_lim=params.upp_lim,
|
|
window=params.window,
|
|
scaler_name='Standard Scaler' if params.use_scaler else 'None',
|
|
scaler_params={} if params.use_scaler else None,
|
|
ar_var=params.target_variable if params.include_ar else None,
|
|
)
|