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)

32
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -865,7 +865,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
prediction_model.fit.assert_called_once_with(pd_merge.return_value)
mlflow_repository.get_prediction_data.assert_called_once_with(
prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable
prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable, 'pyfunc'
)
assert output == {
@@ -951,7 +951,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
prediction_model.fit.assert_called_once_with(transformed_data)
mlflow_repository.get_prediction_data.assert_called_once_with(
prediction_model, transformed_data, 'feat_1'
prediction_model, transformed_data, 'feat_1', 'pyfunc'
)
assert output == {
@@ -1494,8 +1494,11 @@ def test_get_prediction_data_dataframe(mlflow_repository):
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
target_name = 'target'
predict_flavor = 'sklearn'
result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
result = mlflow_repository.get_prediction_data(
prediction_model, retrain_dataset, target_name, predict_flavor
)
prediction_model.predict.assert_called_once_with(retrain_dataset)
assert 'prediction' in result.columns
@@ -1509,11 +1512,32 @@ def test_get_prediction_data_array(mlflow_repository):
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = [5, 6]
target_name = 'target'
predict_flavor = 'sklearn'
result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
result = mlflow_repository.get_prediction_data(
prediction_model, retrain_dataset, target_name, predict_flavor
)
prediction_model.predict.assert_called_once_with(retrain_dataset)
assert 'prediction' in result.columns
assert 'target' in result.columns
assert 'timestamp' in result.columns
assert result.index.tolist() == [0, 1]
def test_get_prediction_data_pyfunc(mlflow_repository):
prediction_model = MagicMock()
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
target_name = 'target'
predict_flavor = 'pyfunc'
result = mlflow_repository.get_prediction_data(
prediction_model, retrain_dataset, target_name, predict_flavor
)
prediction_model.predict.assert_called_once_with({}, retrain_dataset)
assert 'prediction' in result.columns
assert 'target' in result.columns
assert 'timestamp' in result.columns
assert result.index.tolist() == [0, 1]