Files
sientia-dataops-model-manager/model_manager/utils/repository/training_repository.py
Bruno Domingues bee6036205 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.
2025-10-09 15:07:27 -03:00

243 lines
8.9 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.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,
)