lecrapaud 0.18.7__py3-none-any.whl → 0.22.6__py3-none-any.whl

lecrapaud/{api.py → base.py}

@@ -5,24 +5,22 @@ It allows for end-to-end ML workflows including data preprocessing, feature engi
 model training, and prediction.
 
 Basic Usage:
-    # Create a LeCrapaud instance
-    lc = LeCrapaud()
-
     # Create a new experiment
-    experiment = lc.create_experiment(data, target_numbers=[1], target_clf=[1])
+    experiment = LeCrapaud(data=data, target_numbers=[1], target_clf=[1])
 
-    # Train a model
-    best_features, artifacts, best_model = experiment.train(data)
+    # Train the model
+    experiment.fit(data)
 
     # Make predictions
     predictions, scores_reg, scores_clf = experiment.predict(new_data)
 
-    # Or use individual pipeline steps:
-    processed_data = experiment.feature_engineering(data)  # Feature engineering
-    train, val, test = experiment.preprocess_feature(data)  # Data splitting and encoding
-    selected_features = experiment.feature_selection(train)  # Feature selection
-    model_data = experiment.preprocess_model(train, val, test)  # Model preprocessing
-    best_model = experiment.model_selection(model_data)  # Model selection
+    # Load existing experiment
+    experiment = LeCrapaud(id=123)
+    predictions = experiment.predict(new_data)
+
+    # Class methods for experiment management
+    best_exp = LeCrapaud.get_best_experiment_by_name('my_experiment')
+    all_exps = LeCrapaud.list_experiments('my_experiment')
 """
 
 import joblib
@@ -34,84 +32,172 @@ import seaborn as sns
 import numpy as np
 import matplotlib.pyplot as plt
 from lecrapaud.db.session import init_db
-from lecrapaud.feature_selection import FeatureSelectionEngine, PreprocessModel
+from lecrapaud.feature_selection import FeatureSelector
+from lecrapaud.model_preprocessing import ModelPreprocessor
 from lecrapaud.model_selection import (
-    ModelSelectionEngine,
-    ModelEngine,
+    ModelSelector,
+    BaseModel,
     evaluate,
     load_model,
     plot_threshold,
     plot_evaluation_for_classification,
 )
-from lecrapaud.feature_engineering import FeatureEngineeringEngine, PreprocessFeature
+from lecrapaud.feature_engineering import FeatureEngineering
+from lecrapaud.feature_preprocessing import FeaturePreprocessor
 from lecrapaud.experiment import create_experiment
 from lecrapaud.db import Experiment
-from lecrapaud.search_space import normalize_models_idx
+from lecrapaud.search_space import normalize_models_idx, all_models
 from lecrapaud.utils import logger
 from lecrapaud.directories import tmp_dir
 
 
 class LeCrapaud:
-    """Main class for interacting with the LeCrapaud ML pipeline.
+    """
+    Unified LeCrapaud class for machine learning experiments.
 
-    This class provides methods to create and retrieve experiments.
+    This class provides both the ML pipeline functionality and experiment management.
+    It can be initialized either with new data to create an experiment or with an
+    experiment ID to load an existing one.
 
-    Args:
-        uri (str, optional): Database connection URI. If None, uses default connection.
-    """
+    Usage:
+        # Create new experiment
+        experiment = LeCrapaud(data=df, target_numbers=[1, 2], ...)
 
-    def __init__(self, uri: str = None):
-        """Initialize LeCrapaud with optional database URI."""
-        init_db(uri=uri)
+        # Load existing experiment
+        experiment = LeCrapaud(id=123)
 
-    def create_experiment(self, data: pd.DataFrame, **kwargs) -> "ExperimentEngine":
-        """Create a new experiment.
+        # Train the model
+        experiment.fit(data)
 
-        Args:
-            data (pd.DataFrame): Input data for the experiment
-            **kwargs: Additional arguments to configure the experiment
+        # Make predictions
+        predictions = experiment.predict(new_data)
 
-        Returns:
-            ExperimentEngine: A new experiment instance
-        """
-        return ExperimentEngine(data=data, **kwargs)
+    Args:
+        id (int, optional): ID of an existing experiment to load
+        data (pd.DataFrame, optional): Input data for a new experiment
+        uri (str, optional): Database connection URI
+        **kwargs: Additional configuration parameters
+    """
 
-    def get_experiment(self, id: int, **kwargs) -> "ExperimentEngine":
-        """Retrieve an existing experiment by ID.
+    def __init__(
+        self, id: int = None, data: pd.DataFrame = None, uri: str = None, **kwargs
+    ):
+        """Initialize LeCrapaud with either new or existing experiment."""
+        # Initialize database connection
+        init_db(uri=uri)
 
-        Args:
-            id (int): The ID of the experiment to retrieve
-            **kwargs: Additional arguments to pass to the experiment
+        if id:
+            # Load existing experiment
+            self.experiment = Experiment.get(id)
+            # Context from DB takes precedence over kwargs
+            effective_kwargs = {
+                **self.DEFAULT_PARAMS,
+                **kwargs,
+                **self.experiment.context,
+            }
+        else:
+            if data is None:
+                raise ValueError(
+                    "Either id or data must be provided. Data can be a path to a folder containing trained models"
+                )
+            # New experiment: merge defaults with provided kwargs
+            effective_kwargs = {**self.DEFAULT_PARAMS, **kwargs}
 
-        Returns:
-            ExperimentEngine: The retrieved experiment instance
-        """
-        return ExperimentEngine(id=id, **kwargs)
+        # Normalize models_idx if present
+        if "models_idx" in effective_kwargs:
+            effective_kwargs["models_idx"] = normalize_models_idx(
+                effective_kwargs["models_idx"]
+            )
+
+        # Set all parameters as instance attributes
+        for key, value in effective_kwargs.items():
+            setattr(self, key, value)
 
-    def get_last_experiment_by_name(self, name: str, **kwargs) -> "ExperimentEngine":
+        # Create experiment if new
+        if not id:
+            self.experiment = create_experiment(data=data, **effective_kwargs)
+
+        # Create directories
+        experiment_dir = f"{tmp_dir}/{self.experiment.name}"
+        preprocessing_dir = f"{experiment_dir}/preprocessing"
+        data_dir = f"{experiment_dir}/data"
+        os.makedirs(preprocessing_dir, exist_ok=True)
+        os.makedirs(data_dir, exist_ok=True)
+
+    # Default values for all experiment parameters
+    DEFAULT_PARAMS = {
+        # Feature Engineering
+        "columns_drop": [],
+        "columns_boolean": [],
+        "columns_date": [],
+        "columns_te_groupby": [],
+        "columns_te_target": [],
+        "for_training": True,
+        # Preprocessing
+        "time_series": False,
+        "val_size": 0.2,
+        "test_size": 0.2,
+        "columns_pca": [],
+        "pca_temporal": [],
+        "pca_cross_sectional": [],
+        "columns_onehot": [],
+        "columns_binary": [],
+        "columns_ordinal": [],
+        "columns_frequency": [],
+        # Feature Selection
+        "percentile": 20,
+        "corr_threshold": 80,
+        "max_features": 50,
+        "max_p_value_categorical": 0.05,
+        # Model Selection
+        "target_numbers": [],
+        "target_clf": [],
+        "models_idx": [],
+        "max_timesteps": 120,
+        "perform_hyperopt": True,
+        "number_of_trials": 20,
+        "perform_crossval": False,
+        "plot": True,
+        "preserve_model": True,
+        "target_clf_thresholds": {},
+        # Data structure
+        "date_column": None,
+        "group_column": None,
+    }
+
+    @classmethod
+    def get_default_params(cls):
+        """Get the default parameters for experiments."""
+        return cls.DEFAULT_PARAMS.copy()
+
+    def get_effective_context(self):
+        """Get the effective context (merged defaults + experiment context)."""
+        return {k: getattr(self, k, v) for k, v in self.DEFAULT_PARAMS.items()}
+
+    @classmethod
+    def get_last_experiment_by_name(cls, name: str, **kwargs):
         """Retrieve the last experiment by name."""
-        return ExperimentEngine(id=Experiment.get_last_by_name(name).id, **kwargs)
+        return cls(id=Experiment.get_last_by_name(name).id, **kwargs)
 
-    def get_best_experiment_by_name(
-        self, name: str, metric: str = "both", **kwargs
-    ) -> "ExperimentEngine":
+    @classmethod
+    def get_best_experiment_by_name(cls, name: str, **kwargs):
         """Retrieve the best experiment by score."""
-        best_exp = Experiment.get_best_by_score(name=name, metric=metric)
+        best_exp = Experiment.get_best_by_score(name=name)
         if not best_exp:
             return None
-        return ExperimentEngine(id=best_exp.id, **kwargs)
+        return cls(id=best_exp.id, **kwargs)
 
-    def compare_experiment_scores(self, name: str):
-        """Compare scores of experiments with matching names.
+    @classmethod
+    def list_experiments(cls, name: str = None, limit: int = 1000):
+        """List all experiments in the database."""
+        return [
+            cls(id=exp.id) for exp in Experiment.get_all_by_name(name=name, limit=limit)
+        ]
 
-        Args:
-            name (str): Name or partial name of experiments to compare
-
-        Returns:
-            dict: Dictionary containing experiment names as keys and their scores as values
-        """
-        # Get all experiments with the given name pattern
-        experiments = self.list_experiments(name=name)
+    @classmethod
+    def compare_experiment_scores(cls, name: str):
+        """Compare scores of experiments with matching names."""
+        experiments = cls.list_experiments(name=name)
 
         if not experiments:
             return {"error": f"No experiments found with name containing '{name}'"}
@@ -120,9 +206,7 @@ class LeCrapaud:
 
         for exp in experiments:
             for model_sel in exp.experiment.model_selections:
-
                 if model_sel.best_score:
-
                     scores = {
                         "rmse": model_sel.best_score["rmse"],
                         "logloss": model_sel.best_score["logloss"],
@@ -131,7 +215,6 @@ class LeCrapaud:
                         "roc_auc": model_sel.best_score["roc_auc"],
                     }
                     target_name = model_sel.target.name
-
                     comparison[exp.experiment.name][target_name] = scores
                 else:
                     logger.warning(
@@ -140,71 +223,143 @@ class LeCrapaud:
 
         return comparison
 
-    def list_experiments(
-        self, name: str = None, limit: int = 1000
-    ) -> list["ExperimentEngine"]:
-        """List all experiments in the database."""
-        return [
-            ExperimentEngine(id=exp.id)
-            for exp in Experiment.get_all_by_name(name=name, limit=limit)
-        ]
-
-
-class ExperimentEngine:
-    """Engine for managing ML experiments.
-
-    This class handles the complete ML pipeline including feature engineering,
-    model training, and prediction. It can be initialized with either new data
-    or by loading an existing experiment by ID.
-
-    Args:
-        id (int, optional): ID of an existing experiment to load
-        data (pd.DataFrame, optional): Input data for a new experiment
-        **kwargs: Additional configuration parameters
-    """
+    # Main ML Pipeline Methods
+    # ========================
 
-    def __init__(self, id: int = None, data: pd.DataFrame = None, **kwargs):
-        """Initialize the experiment engine with either new or existing experiment."""
-        if id:
-            self.experiment = Experiment.get(id)
-            kwargs.update(self.experiment.context)
-            experiment_dir = f"{tmp_dir}/{self.experiment.name}"
-            preprocessing_dir = f"{experiment_dir}/preprocessing"
-            data_dir = f"{experiment_dir}/data"
-            os.makedirs(preprocessing_dir, exist_ok=True)
-            os.makedirs(data_dir, exist_ok=True)
-        else:
-            if data is None:
-                raise ValueError(
-                    "Either id or data must be provided. Data can be a path to a folder containing trained models"
-                )
-            self.experiment = create_experiment(data=data, **kwargs)
+    def fit(self, data, best_params=None):
+        """
+        Fit the complete ML pipeline on the provided data.
 
-        # Set all kwargs as instance attributes
-        for key, value in kwargs.items():
-            if key == "models_idx":
-                value = normalize_models_idx(value)
-            setattr(self, key, value)
+        Args:
+            data (pd.DataFrame): Input training data
+            best_params (dict, optional): Pre-defined best parameters
 
-    def train(self, data, best_params=None):
+        Returns:
+            self: Returns self for chaining
+        """
         logger.info("Running training...")
 
-        data_eng = self.feature_engineering(data)
+        # Step 1: Feature Engineering
+        logger.info("Starting feature engineering...")
+        feature_eng = FeatureEngineering(experiment=self.experiment)
+        feature_eng.fit(data)
+        data_eng = feature_eng.get_data()
         logger.info("Feature engineering done.")
 
-        train, val, test = self.preprocess_feature(data_eng)
+        # Step 2: Feature Preprocessing (split data)
+        logger.info("Starting feature preprocessing...")
+        from lecrapaud.feature_preprocessing import split_data
+
+        train, val, test = split_data(data_eng, experiment=self.experiment)
+
+        # Apply feature preprocessing transformations
+        feature_preprocessor = FeaturePreprocessor(experiment=self.experiment)
+        feature_preprocessor.fit(train)
+        train = feature_preprocessor.transform(train)
+        if val is not None:
+            val = feature_preprocessor.transform(val)
+        if test is not None:
+            test = feature_preprocessor.transform(test)
         logger.info("Feature preprocessing done.")
 
-        self.feature_selection(train)
+        # Step 3: Feature Selection (for each target)
+        logger.info("Starting feature selection...")
+        for target_number in self.target_numbers:
+            feature_selector = FeatureSelector(
+                experiment=self.experiment, target_number=target_number
+            )
+            feature_selector.fit(train)
+
+        # Refresh experiment to get updated features
+        self.experiment = Experiment.get(self.experiment.id)
+        all_features = self.experiment.get_all_features(
+            date_column=self.date_column, group_column=self.group_column
+        )
+        joblib.dump(
+            all_features, f"{self.experiment.path}/preprocessing/all_features.pkl"
+        )
         logger.info("Feature selection done.")
 
-        std_data, reshaped_data = self.preprocess_model(train, val, test)
+        # Step 4: Model Preprocessing (scaling)
+        logger.info("Starting model preprocessing...")
+        model_preprocessor = ModelPreprocessor(experiment=self.experiment)
+
+        # Fit and transform training data, then transform val/test
+        model_preprocessor.fit(train)
+        train_scaled = model_preprocessor.transform(train)
+        val_scaled = model_preprocessor.transform(val) if val is not None else None
+        test_scaled = model_preprocessor.transform(test) if test is not None else None
+
+        # Create data dict for model selection (keep both raw and scaled splits)
+        std_data = {
+            "train": train,
+            "val": val,
+            "test": test,
+            "train_scaled": train_scaled,
+            "val_scaled": val_scaled,
+            "test_scaled": test_scaled,
+        }
+
+        # Handle time series reshaping if needed
+        reshaped_data = None
+        # Check if any model requires recurrent processing
+        need_reshaping = (
+            any(all_models[i].get("recurrent") for i in self.models_idx)
+            and self.time_series
+        )
+
+        if need_reshaping:
+            # Sanity check: make sure we have enough data for max_timesteps
+            if (
+                self.group_column
+                and train_scaled.groupby(self.group_column).size().min()
+                < self.max_timesteps
+            ) or train_scaled.shape[0] < self.max_timesteps:
+                raise ValueError(
+                    f"Not enough data for group_column {self.group_column} to reshape data for recurrent models"
+                )
+
+            from lecrapaud.model_preprocessing import reshape_time_series
+
+            features = self.experiment.get_all_features(
+                date_column=self.date_column, group_column=self.group_column
+            )
+            reshaped_data = reshape_time_series(
+                self.experiment,
+                features,
+                train_scaled,
+                val_scaled,
+                test_scaled,
+                timesteps=self.max_timesteps,
+            )
         logger.info("Model preprocessing done.")
 
-        self.model_selection(std_data, reshaped_data, best_params=best_params)
+        # Step 5: Model Selection (for each target)
+        logger.info("Starting model selection...")
+        self.models_ = {}
+        for target_number in self.target_numbers:
+            model_selector = ModelSelector(
+                experiment=self.experiment, target_number=target_number
+            )
+            model_selector.fit(
+                std_data, reshaped_data=reshaped_data, best_params=best_params
+            )
+            self.models_[target_number] = model_selector.get_best_model()
         logger.info("Model selection done.")
 
+        return self
+
     def predict(self, new_data, verbose: int = 0):
+        """
+        Make predictions on new data using the trained pipeline.
+
+        Args:
+            new_data (pd.DataFrame): Input data for prediction
+            verbose (int): Verbosity level (0=warnings only, 1=all logs)
+
+        Returns:
+            tuple: (predictions_df, scores_regression, scores_classification)
+        """
         # for scores if TARGET is in columns
         scores_reg = []
         scores_clf = []
@@ -214,28 +369,67 @@ class ExperimentEngine:
 
         logger.warning("Running prediction...")
 
-        # feature engineering + preprocessing
-        data = self.feature_engineering(
-            data=new_data,
-            for_training=False,
-        )
-        data = self.preprocess_feature(data, for_training=False)
-        data, scaled_data, reshaped_data = self.preprocess_model(
-            data, for_training=False
+        # Apply the same preprocessing pipeline as training
+        # Step 1: Feature Engineering
+        feature_eng = FeatureEngineering(experiment=self.experiment)
+        feature_eng.fit(new_data)
+        data = feature_eng.get_data()
+
+        # Step 2: Feature Preprocessing (no splitting for prediction)
+        feature_preprocessor = FeaturePreprocessor(experiment=self.experiment)
+        # Load existing transformations and apply
+        data = feature_preprocessor.transform(data)
+
+        # Step 3: Model Preprocessing (scaling)
+        model_preprocessor = ModelPreprocessor(experiment=self.experiment)
+        # Apply existing scaling
+        scaled_data = model_preprocessor.transform(data)
+
+        # Step 4: Time series reshaping if needed
+        reshaped_data = None
+        # Check if any model requires recurrent processing
+        need_reshaping = (
+            any(all_models[i].get("recurrent") for i in self.models_idx)
+            and self.time_series
         )
 
-        for target_number in self.target_numbers:
+        if need_reshaping:
+            # Sanity check: make sure we have enough data for max_timesteps
+            if (
+                self.group_column
+                and scaled_data.groupby(self.group_column).size().min()
+                < self.max_timesteps
+            ) or scaled_data.shape[0] < self.max_timesteps:
+                raise ValueError(
+                    f"Not enough data for group_column {self.group_column} to reshape data for recurrent models"
+                )
 
-            # loading model
+            from lecrapaud.model_preprocessing import reshape_time_series
+
+            all_features = self.experiment.get_all_features(
+                date_column=self.date_column, group_column=self.group_column
+            )
+            # For prediction, we reshape the entire dataset
+            reshaped_data = reshape_time_series(
+                self.experiment, all_features, scaled_data, timesteps=self.max_timesteps
+            )
+            reshaped_data = reshaped_data[
+                "x_train_reshaped"
+            ]  # Only need X data for prediction
+
+        # Step 5: Predict for each target
+        for target_number in self.target_numbers:
+            # Load the trained model
             target_dir = f"{self.experiment.path}/TARGET_{target_number}"
+            model = BaseModel(path=target_dir, target_number=target_number)
+
+            # Get features for this target
             all_features = self.experiment.get_all_features(
                 date_column=self.date_column, group_column=self.group_column
             )
             features = self.experiment.get_features(target_number)
 
-            model = ModelEngine(path=target_dir, target_number=target_number)
-
-            # getting data
+            # Prepare prediction data
             if model.recurrent:
                 features_idx = [
                     i for i, e in enumerate(all_features) if e in set(features)
@@ -244,16 +438,14 @@ class ExperimentEngine:
             else:
                 x_pred = scaled_data[features] if model.need_scaling else data[features]
 
-            # predicting
+            # Make prediction
             y_pred = model.predict(x_pred)
 
-            # fix for recurrent model because x_val has no index as it is a 3D np array
+            # Fix index for recurrent models
             if model.recurrent:
-                y_pred.index = (
-                    new_data.index
-                )  # TODO: not sure this will work for old experiment not aligned with data_for_training for test use case (done, this is why we decode the test set)
+                y_pred.index = new_data.index
 
-            # unscaling prediction
+            # Unscale prediction if needed
             if (
                 model.need_scaling
                 and model.target_type == "regression"
@@ -267,7 +459,7 @@ class ExperimentEngine:
                 )
                 y_pred.name = "PRED"
 
-            # evaluate if TARGET is in columns (case-insensitive check)
+            # Evaluate if target is present in new_data
             target_col = next(
                 (
                     col
@@ -291,126 +483,21 @@ class ExperimentEngine:
                 else:
                     scores_reg.append(score)
 
-            # renaming and concatenating with initial data
+            # Add predictions to the output dataframe
             if isinstance(y_pred, pd.DataFrame):
                 y_pred = y_pred.add_prefix(f"TARGET_{target_number}_")
                 new_data = pd.concat([new_data, y_pred], axis=1)
-
             else:
                 y_pred.name = f"TARGET_{target_number}_PRED"
                 new_data = pd.concat([new_data, y_pred], axis=1)
 
+        # Format scores
         if len(scores_reg) > 0:
             scores_reg = pd.DataFrame(scores_reg).set_index("TARGET")
         if len(scores_clf) > 0:
             scores_clf = pd.DataFrame(scores_clf).set_index("TARGET")
-        return new_data, scores_reg, scores_clf
-
-    def feature_engineering(self, data, for_training=True):
-        app = FeatureEngineeringEngine(
-            data=data,
-            columns_drop=getattr(self, "columns_drop", []),
-            columns_boolean=getattr(self, "columns_boolean", []),
-            columns_date=getattr(self, "columns_date", []),
-            columns_te_groupby=getattr(self, "columns_te_groupby", []),
-            columns_te_target=getattr(self, "columns_te_target", []),
-            for_training=getattr(self, "for_training", True),
-        )
-        data = app.run()
-        return data
-
-    def preprocess_feature(self, data, for_training=True):
-        app = PreprocessFeature(
-            data=data,
-            experiment=getattr(self, "experiment", None),
-            time_series=getattr(self, "time_series", False),
-            date_column=getattr(self, "date_column", None),
-            group_column=getattr(self, "group_column", None),
-            val_size=getattr(self, "val_size", 0.2),
-            test_size=getattr(self, "test_size", 0.2),
-            columns_pca=getattr(self, "columns_pca", []),
-            pca_temporal=getattr(self, "pca_temporal", []),
-            pca_cross_sectional=getattr(self, "pca_cross_sectional", []),
-            columns_onehot=getattr(self, "columns_onehot", []),
-            columns_binary=getattr(self, "columns_binary", []),
-            columns_ordinal=getattr(self, "columns_ordinal", []),
-            columns_frequency=getattr(self, "columns_frequency", []),
-            target_numbers=getattr(self, "target_numbers", []),
-            target_clf=getattr(self, "target_clf", []),
-        )
-        if for_training:
-            train, val, test = app.run()
-            return train, val, test
-        else:
-            data = app.inference()
-            return data
 
-    def feature_selection(self, train):
-        for target_number in self.target_numbers:
-            app = FeatureSelectionEngine(
-                train=train,
-                target_number=target_number,
-                experiment=self.experiment,
-                target_clf=self.target_clf,
-            )
-            app.run()
-        self.experiment = Experiment.get(self.experiment.id)
-        all_features = self.experiment.get_all_features(
-            date_column=self.date_column, group_column=self.group_column
-        )
-        joblib.dump(
-            all_features, f"{self.experiment.path}/preprocessing/all_features.pkl"
-        )
-        return all_features
-
-    def preprocess_model(self, train, val=None, test=None, for_training=True):
-        app = PreprocessModel(
-            train=train,
-            val=val,
-            test=test,
-            experiment=getattr(self, "experiment", None),
-            target_numbers=getattr(self, "target_numbers", []),
-            target_clf=getattr(self, "target_clf", []),
-            models_idx=getattr(self, "models_idx", []),
-            time_series=getattr(self, "time_series", False),
-            max_timesteps=getattr(self, "max_timesteps", 120),
-            date_column=getattr(self, "date_column", None),
-            group_column=getattr(self, "group_column", None),
-        )
-        if for_training:
-            data, reshaped_data = app.run()
-            return data, reshaped_data
-        else:
-            data, scaled_data, reshaped_data = app.inference()
-            return data, scaled_data, reshaped_data
-
-    def model_selection(self, data, reshaped_data, best_params=None):
-        for target_number in self.target_numbers:
-            app = ModelSelectionEngine(
-                data=data,
-                reshaped_data=reshaped_data,
-                target_number=target_number,
-                experiment=getattr(self, "experiment", None),
-                target_clf=getattr(self, "target_clf", []),
-                models_idx=getattr(self, "models_idx", []),
-                time_series=getattr(self, "time_series", False),
-                date_column=getattr(self, "date_column", None),
-                group_column=getattr(self, "group_column", None),
-                target_clf_thresholds=getattr(self, "target_clf_thresholds", {}),
-            )
-            if best_params and target_number not in best_params.keys():
-                raise ValueError(
-                    f"Target {target_number} not found in best_params passed as argument"
-                )
-            app.run(
-                self.experiment_name,
-                perform_hyperopt=self.perform_hyperopt,
-                number_of_trials=self.number_of_trials,
-                perform_crossval=self.perform_crossval,
-                plot=self.plot,
-                preserve_model=self.preserve_model,
-                best_params=best_params[target_number] if best_params else None,
-            )
+        return new_data, scores_reg, scores_clf
 
     def get_scores(self, target_number: int):
         return pd.read_csv(
@@ -475,6 +562,9 @@ class ExperimentEngine:
             # For lightgbm models
             importances = model.feature_importance(importance_type="split")
             importance_type = "Split"
+        elif hasattr(model, "get_feature_importance"):
+            importances = model.get_feature_importance()
+            importance_type = "Feature importance"
         elif hasattr(model, "coef_"):
             # For linear models
             importances = np.abs(model.coef_.flatten())
@@ -557,7 +647,7 @@ class ExperimentEngine:
         if not os.path.exists(params_file):
             raise FileNotFoundError(
                 f"Best parameters file not found at {params_file}. "
-                "Make sure to run model training first."
+                "Make sure to fit model training first."
             )
 
         try: