""" 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') 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) 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', ar_var=params.target_variable if params.include_ar else None, )