SIENTIAPDE-1251: Implement ML model training activity and repository
This commit introduces the 'Training' activity and 'TrainingRepository' for handling ML model training operations within the Model Manager system. - Added model_manager/activities/training.py for the Training activity, which extends BaseActivity and integrates with Temporal workflows. - Added model_manager/utils/repository/training_repository.py for the TrainingRepository, which encapsulates the core training logic. - Updated model_manager/activities/activities.py to include the Training activity in the main activities orchestrator. - Updated README.md to document the new 'Training' component. - Added unit tests for the new activity and repository.
This commit is contained in:
Bruno Domingues
242
model_manager/utils/repository/training_repository.py
Normal file
242
model_manager/utils/repository/training_repository.py
Normal file
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
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.linear_models import LinearRegressionModel
|
||||
from sientia.metrics import mae, mse, r2
|
||||
from sientia.preprocessing import DataPreprocessor
|
||||
from sientia.utils import split_train_test
|
||||
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.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
|
||||
"""
|
||||
# Load data from BytesIO
|
||||
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
|
||||
|
||||
# Initialize and fit data preprocessor
|
||||
process_data = self.init_data_preprocessor(params)
|
||||
process_data.fit(data)
|
||||
data_view = process_data.transform(data)
|
||||
|
||||
# Validate transformed data
|
||||
if len(data_view) <= 0:
|
||||
raise ValueError('Data view is empty after transformation')
|
||||
|
||||
# Split 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,
|
||||
)
|
||||
|
||||
# Prepare training data
|
||||
data_train = pd.concat([x_train, y_train], axis=1)
|
||||
scaler_dict = self.init_scaler_dict(process_data, params)
|
||||
|
||||
# Create and train linear regression model
|
||||
regr = LinearRegressionModel(
|
||||
target_variable=params.target_variable,
|
||||
variable_columns=params.variable_columns,
|
||||
)
|
||||
regr.fit(data_train)
|
||||
|
||||
# Return training result
|
||||
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 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 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)
|
||||
"""
|
||||
# Make predictions on test set
|
||||
tmr.y_pred = tmr.regr.predict(tmr.x_test)
|
||||
|
||||
# Denormalize data if scaler was used
|
||||
if params.use_scaler:
|
||||
scaler = tmr.process_data.get_scaler()
|
||||
|
||||
# Denormalize 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)
|
||||
|
||||
# Denormalize 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)
|
||||
tmr.y_pred = scaler.denormalize_predictions(tmr.y_pred, params.target_variable)
|
||||
|
||||
# Add index to predictions
|
||||
tmr.y_pred = pd.Series(tmr.y_pred, index=tmr.y_test.index)
|
||||
tmr.y_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
# Reorder 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()
|
||||
|
||||
# Calculate evaluation metrics
|
||||
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_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,
|
||||
)
|
||||
Reference in New Issue
Block a user