update model auto trainer

pycaret_util.py

@@ -3,9 +3,9 @@ from loguru import logger
 
 
 class GetBestModel:
-    def __init__(self, df, target_col_list, list_of_models, no_of_models=1):
+    def __init__(self, df, target_col, list_of_models, no_of_models=1):
         self.df = df
-        self.target_col_list = target_col_list
+        self.target_col = target_col
         self.no_of_models = no_of_models
         self.list_of_models = list_of_models
 
@@ -18,7 +18,7 @@ class GetBestModel:
         """
         try:
             logger.info("Using Pycaret to train mentioned models")
-            regression.setup(data=self.df, target=self.target_col_list)
+            regression.setup(data=self.df, target=self.target_col)
             logger.info(f"Selecting the best model using the metric {comparison_metric}")
             best_model = regression.compare_models(include=self.list_of_models, sort=comparison_metric,
                                                    n_select=self.no_of_models)