""" 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 """ data = load_data(uploaded_file, params.line_separator, params.decimal_separator) 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') 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, ) data_train = pd.concat([x_train, y_train], axis=1) scaler_dict = self._init_scaler_dict(process_data, params) regr = LinearRegressionModel( target_variable=params.target_variable, variable_columns=params.variable_columns, ) regr.fit(data_train) self.logger.info( f'Model trained successfully - experiment run id: {params.experiment_run_id}' ) 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) """ y_pred_array = tmr.regr.predict(tmr.x_test) if params.use_scaler: scaler = tmr.process_data.get_scaler() # 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) 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 tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index) tmr.y_pred.name = f'{params.target_variable}_pred' 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() 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) # Extract model equation tmr.equation = self._extract_model_equation(tmr.regr, params) self.logger.info( f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}' ) 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, ) def _extract_model_equation( self, regr: LinearRegressionModel, params: TrainModelParams ) -> dict: """ Extract the linear regression equation coefficients and create equation metadata. This method extracts the coefficients and intercept from the trained model and creates a structured dictionary containing the equation information for serialization as JSON artifact. Args: regr: Trained LinearRegressionModel object params: Training parameters containing variable information Returns: dict: Equation metadata containing: - target_variable: Name of the target variable - coefficients: Dictionary mapping variable names to coefficients - intercept: Model intercept value - equation_string: Human-readable equation string - latex_equation: LaTeX formatted equation """ coefficients = regr.regr.coef_ intercept = regr.regr.intercept_ # Create coefficients dictionary coefficients_dict = {} for i, var in enumerate(params.variable_columns): coefficients_dict[var] = float(coefficients[i]) # Create equation string equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()] equation_string = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join( equation_parts ) # Create LaTeX equation latex_parts = [f'{coef:.6f} \\cdot {var}' for var, coef in coefficients_dict.items()] latex_equation = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(latex_parts) return { 'target_variable': params.target_variable, 'coefficients': coefficients_dict, 'intercept': float(intercept), 'equation_string': equation_string, 'latex_equation': latex_equation, 'model_type': 'Linear Regression', }