SIENTIAPDE-1171

Update GITHUB_BRANCH in values.yaml and refactor MLFlowRepository for model retraining

- Changed GITHUB_BRANCH in values.yaml to reflect the new pipeline for retraining.
- Refactored MLFlowRepository by introducing create_model_experiment and perform_model_retrain methods to streamline model retraining processes.
- Updated test cases to cover new functionalities and ensure proper error handling in model management workflows.

laborious/utils/repository/model_repository.py

@@ -116,32 +116,7 @@ class MLFlowRepository():
         next_run_number = len(runs) + 1
         return f"{model_name}-{next_run_number}"
 
-    def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
-        """
-        Retrain a model with new data.
-
-        Parameters:
-            data (pandas.DataFrame): The new data to use for retraining.
-            model_name (str): The name of the model to retrain.
-            metrics_list (list): The metrics to be used to compare the models.
-            compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one.
-                                If False, the retrain will always be considered.
-            split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets.
-                If False, the data will be used as a unique block for retraining.
-            update_report (bool): If True, a report will be created with the data of the retrained model.
-            update_transformation (bool): If True, the model will be updated in the MLflow tracking server.
-            update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server.
-            shuffle_data (bool): If True, the data will be shuffled before splitting.
-            model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'.
-
-
-        Returns:
-            mlflow.sklearn.Model: The retrained prediction model.
-            mlflow.sklearn.Model: The retrained data model.
-            mse (float): The mean squared error of the retrained model.
-            r2 (float): The R-squared score of the retrained model.
-        """
-
+    def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
         # load predictor model
         predictor_uri = f"models:/{model_name}/production"
         # load transform model
@@ -156,18 +131,27 @@ class MLFlowRepository():
         prediction_model = mlflow.sklearn.load_model(predictor_uri)
         data_model = data_model.fit(data)
         treated_data = data_model.predict(data)
-        # align target column with treated_data
+
         target_name = data_model.target_variable
         y = data[target_name]
         treated_data = pd.merge(
             treated_data, y, left_index=True, right_index=True)
         prediction_model = prediction_model.fit(treated_data)
-        # Example usage
         experiment = self.get_experiment_by_run_id(latest_production_id)
+        mlflow.set_experiment(experiment)
+
+        return prediction_model, data_model, experiment
+
+    def perform_model_retrain(self,
+                              prediction_model,
+                              data_model,
+                              experiment: str,
+                              model_name: str,
+                              data: pd.DataFrame):
+
         pred_model_atributes = vars(prediction_model)  # load class attributes
         data_model_atributes = vars(data_model)  # load class attributes
-        mlflow.set_experiment(experiment)
-        experiment_description = "Retrain model {model_name} with new data"
+        experiment_description = f"Retrain model {model_name} with new data"
         current_run_name = self.get_next_run_name(experiment)
         with mlflow.start_run(
             run_name=current_run_name, description=experiment_description
@@ -203,6 +187,14 @@ class MLFlowRepository():
 
             return "Model retrained successfully", experiment
 
+    def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
+
+        prediction_model, data_model, experiment = self.create_model_experiment(
+            model_name, data)
+        retrain_result = self.perform_model_retrain(
+            prediction_model, data_model, experiment, model_name, data)
+        return retrain_result
+
     def get_experiment(self, experiment_name: str) -> int:
         experiment = mlflow.get_experiment_by_name(experiment_name)