gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1273

Browse Source 
Enhance prediction data retrieval in MLFlowRepository by adding predict_flavor parameter

- Updated get_prediction_data method to accept a predict_flavor argument, allowing for different prediction model handling.
- Adjusted calls to get_prediction_data throughout the codebase to include the new parameter.
- Added new test cases to validate behavior for different predict_flavor values, ensuring robust functionality.
This commit is contained in:
vitor-aignosi
2025-11-19 13:26:08 -03:00
parent 8f6a3cd93d
commit 684bc9fb31
2 changed files with 34 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
laborious/utils/repository/model_repository.py
View File

@@ -705,14 +705,17 @@ class MLFlowRepository(SientiaMonitoring):
"""
def get_prediction_data(
self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str
self, prediction_model: Any, retrain_dataset: pd.DataFrame, target_name: str, predict_flavor: str
) -> pd.DataFrame:
"""
Get prediction data from prediction model.
"""
input_index = retrain_dataset.index
prediction_data = prediction_model.predict(retrain_dataset)
if predict_flavor == 'pyfunc':
prediction_data = prediction_model.predict({}, retrain_dataset)
else:
prediction_data = prediction_model.predict(retrain_dataset)
if isinstance(prediction_data, pd.DataFrame):
prediction_data.columns = pd.Index(['prediction'])
@@ -858,7 +861,7 @@ class MLFlowRepository(SientiaMonitoring):
prediction_model.fit(retrain_dataset)
# get prediction data
prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name)
prediction_data = self.get_prediction_data(prediction_model, retrain_dataset, target_name, predict_flavor)
self.info(f'Model experiment creation completed successfully for {model_name}', metadata)