upgini 1.2.114a4__py3-none-any.whl → 1.2.115a1__py3-none-any.whl

upgini/features_enricher.py

@@ -1,8 +1,6 @@
 import dataclasses
 import datetime
 import gc
-import hashlib
-import itertools
 import json
 import logging
 import os
@@ -14,7 +12,7 @@ from collections import Counter
 from copy import deepcopy
 from dataclasses import dataclass
 from threading import Thread
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable
 
 import numpy as np
 import pandas as pd
@@ -101,7 +99,7 @@ from upgini.utils.email_utils import (
 from upgini.utils.feature_info import FeatureInfo, _round_shap_value
 from upgini.utils.features_validator import FeaturesValidator
 from upgini.utils.format import Format
-from upgini.utils.hash_utils import file_hash
+from upgini.utils.hash_utils import file_hash, hash_input
 from upgini.utils.ip_utils import IpSearchKeyConverter
 from upgini.utils.phone_utils import PhoneSearchKeyDetector
 from upgini.utils.postal_code_utils import PostalCodeSearchKeyDetector
@@ -113,10 +111,11 @@ except Exception:
         CustomFallbackProgressBar as ProgressBar,
     )
 
+from upgini.utils.config import SampleConfig
 from upgini.utils.psi import calculate_features_psi, calculate_sparsity_psi
-from upgini.utils.sample_utils import SampleColumns, SampleConfig, _num_samples, sample
+from upgini.utils.sample_utils import SampleColumns, _num_samples, sample
 from upgini.utils.sort import sort_columns
-from upgini.utils.target_utils import calculate_psi, define_task
+from upgini.utils.target_utils import calculate_psi, define_task, is_imbalanced
 from upgini.utils.warning_counter import WarningCounter
 from upgini.version_validator import validate_version
 
@@ -132,7 +131,7 @@ class FeaturesEnricher(TransformerMixin):
     Parameters
     ----------
     search_keys: dict of str->SearchKey or int->SearchKey
-        Dictionary with column names or indices mapping to key types.
+        dictionary with column names or indices mapping to key types.
         Each of this columns will be used as a search key to find features.
 
     country_code: str, optional (default=None)
@@ -164,7 +163,7 @@ class FeaturesEnricher(TransformerMixin):
         Custom loss function to use for feature selection and metrics calculation.
 
     shared_datasets: list of str, optional (default=None)
-        List of private shared dataset ids for custom search
+        list of private shared dataset ids for custom search
     """
 
     TARGET_NAME = "target"
@@ -208,34 +207,34 @@ class FeaturesEnricher(TransformerMixin):
 
     def __init__(
         self,
-        search_keys: Optional[Dict[str, SearchKey]] = None,
-        country_code: Optional[str] = None,
-        model_task_type: Optional[Union[ModelTaskType, str]] = None,
-        api_key: Optional[str] = None,
-        endpoint: Optional[str] = None,
-        search_id: Optional[str] = None,
-        shared_datasets: Optional[List[str]] = None,
-        runtime_parameters: Optional[RuntimeParameters] = None,
-        date_format: Optional[str] = None,
+        search_keys: dict[str, SearchKey] | None = None,
+        country_code: str | None = None,
+        model_task_type: ModelTaskType | str | None = None,
+        api_key: str | None = None,
+        endpoint: str | None = None,
+        search_id: str | None = None,
+        shared_datasets: list[str] | None = None,
+        runtime_parameters: RuntimeParameters | None = None,
+        date_format: str | None = None,
         random_state: int = 42,
-        cv: Optional[CVType] = None,
-        loss: Optional[str] = None,
+        cv: CVType | None = None,
+        loss: str | None = None,
         autodetect_search_keys: bool = True,
-        generate_features: Optional[List[str]] = None,
-        columns_for_online_api: Optional[List[str]] = None,
-        round_embeddings: Optional[int] = None,
+        generate_features: list[str] | None = None,
+        columns_for_online_api: list[str] | None = None,
+        round_embeddings: int | None = None,
         logs_enabled: bool = True,
         raise_validation_error: bool = True,
-        exclude_columns: Optional[List[str]] = None,
-        baseline_score_column: Optional[Any] = None,
-        client_ip: Optional[str] = None,
-        client_visitorid: Optional[str] = None,
-        custom_bundle_config: Optional[str] = None,
+        exclude_columns: list[str] | None = None,
+        baseline_score_column: Any | None = None,
+        client_ip: str | None = None,
+        client_visitorid: str | None = None,
+        custom_bundle_config: str | None = None,
         add_date_if_missing: bool = True,
         disable_force_downsampling: bool = False,
-        id_columns: Optional[List[str]] = None,
+        id_columns: list[str] | None = None,
         generate_search_key_features: bool = True,
-        sample_config: Optional[SampleConfig] = None,
+        sample_config: SampleConfig | None = None,
         print_trace_id: bool = False,
         **kwargs,
     ):
@@ -259,21 +258,21 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(msg)
             print(msg)
 
-        self.passed_features: List[str] = []
-        self.df_with_original_index: Optional[pd.DataFrame] = None
-        self.fit_columns_renaming: Optional[Dict[str, str]] = None
+        self.passed_features: list[str] = []
+        self.df_with_original_index: pd.DataFrame | None = None
+        self.fit_columns_renaming: dict[str, str] | None = None
         self.country_added = False
-        self.fit_generated_features: List[str] = []
-        self.fit_dropped_features: Set[str] = set()
+        self.fit_generated_features: list[str] = []
+        self.fit_dropped_features: set[str] = set()
         self.fit_search_keys = search_keys
         self.warning_counter = WarningCounter()
-        self.X: Optional[pd.DataFrame] = None
-        self.y: Optional[pd.Series] = None
-        self.eval_set: Optional[List[Tuple]] = None
-        self.autodetected_search_keys: Dict[str, SearchKey] = {}
+        self.X: pd.DataFrame | None = None
+        self.y: pd.Series | None = None
+        self.eval_set: list[tuple] | None = None
+        self.autodetected_search_keys: dict[str, SearchKey] = {}
         self.imbalanced = False
         self.fit_select_features = True
-        self.__cached_sampled_datasets: Dict[str, Tuple[pd.DataFrame, pd.DataFrame, pd.Series, Dict, Dict, Dict]] = (
+        self.__cached_sampled_datasets: dict[str, tuple[pd.DataFrame, pd.DataFrame, pd.Series, dict, dict, dict]] = (
             dict()
         )
 
@@ -289,19 +288,17 @@ class FeaturesEnricher(TransformerMixin):
             self.model_task_type = ModelTaskType.parse(model_task_type)
         self.model_task_type = model_task_type
         self.endpoint = endpoint
-        self._search_task: Optional[SearchTask] = None
+        self._search_task: SearchTask | None = None
         self.features_info: pd.DataFrame = self.EMPTY_FEATURES_INFO
         self._features_info_without_links: pd.DataFrame = self.EMPTY_FEATURES_INFO
         self._internal_features_info: pd.DataFrame = self.EMPTY_INTERNAL_FEATURES_INFO
         self.relevant_data_sources: pd.DataFrame = self.EMPTY_DATA_SOURCES
         self._relevant_data_sources_wo_links: pd.DataFrame = self.EMPTY_DATA_SOURCES
-        self.metrics: Optional[pd.DataFrame] = None
+        self.metrics: pd.DataFrame | None = None
         self.feature_names_ = []
         self.external_source_feature_names = []
-        self.zero_shap_client_features = []
-        self.unstable_client_features = []
         self.feature_importances_ = []
-        self.psi_values: Optional[Dict[str, float]] = None
+        self.psi_values: dict[str, float] | None = None
         self.search_id = search_id
         self.disable_force_downsampling = disable_force_downsampling
         self.print_trace_id = print_trace_id
@@ -321,7 +318,8 @@ class FeaturesEnricher(TransformerMixin):
                     x_columns = [c.name for c in file_metadata.columns]
                     self.fit_columns_renaming = {c.name: c.originalName for c in file_metadata.columns}
                     df = pd.DataFrame(columns=x_columns)
-                    self.__prepare_feature_importances(trace_id, df, silent=True)
+                    self.__prepare_feature_importances(trace_id, df, silent=True, update_selected_features=False)
+                    self.__show_selected_features()
                     # TODO validate search_keys with search_keys from file_metadata
                     print(self.bundle.get("search_by_task_id_finish"))
                     self.logger.debug(f"Successfully initialized with search_id: {search_id}")
@@ -377,7 +375,7 @@ class FeaturesEnricher(TransformerMixin):
         self.autofe_features_display_handle = None
         self.report_button_handle = None
 
-    def _get_sample_config(self, sample_config: Optional[SampleConfig] = None):
+    def _get_sample_config(self, sample_config: SampleConfig | None = None):
         sample_config = sample_config or SampleConfig(force_sample_size=Dataset.FORCE_SAMPLE_SIZE)
 
         maybe_downsampling_limit = self.runtime_parameters.properties.get("downsampling_limit")
@@ -432,22 +430,21 @@ class FeaturesEnricher(TransformerMixin):
 
     def fit(
         self,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray],
-        y: Union[pd.Series, np.ndarray, List],
-        eval_set: Optional[Union[List[tuple], tuple]] = None,
+        X: pd.DataFrame | pd.Series | np.ndarray,
+        y: pd.Series | np.ndarray | list,
+        eval_set: list[tuple] | tuple | None = None,
         *args,
-        exclude_features_sources: Optional[List[str]] = None,
-        calculate_metrics: Optional[bool] = None,
-        estimator: Optional[Any] = None,
-        scoring: Union[Callable, str, None] = None,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
-        remove_outliers_calc_metrics: Optional[bool] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
-        search_id_callback: Optional[Callable[[str], Any]] = None,
+        exclude_features_sources: list[str] | None = None,
+        calculate_metrics: bool | None = None,
+        estimator: Any | None = None,
+        scoring: Callable | str | None = None,
+        remove_outliers_calc_metrics: bool | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
+        search_id_callback: Callable[[str], Any] | None = None,
         select_features: bool = True,
-        auto_fe_parameters: Optional[AutoFEParameters] = None,
-        stability_threshold: float = 0.15,
+        auto_fe_parameters: AutoFEParameters | None = None,
+        stability_threshold: float = 0.2,
+        stability_agg_func: str = "max",
         **kwargs,
     ):
         """Fit to data.
@@ -462,14 +459,8 @@ class FeaturesEnricher(TransformerMixin):
         y: array-like of shape (n_samples,)
             Target values.
 
-        eval_set: List[tuple], optional (default=None)
-            List of pairs (X, y) for validation.
-
-        importance_threshold: float, optional (default=None)
-            Minimum SHAP value to select a feature. Default value is 0.0.
-
-        max_features: int, optional (default=None)
-            Maximum number of most important features to select. If None, the number is unlimited.
+        eval_set: list[tuple], optional (default=None)
+            list of pairs (X, y) for validation.
 
         calculate_metrics: bool, optional (default=None)
             Whether to calculate and show metrics.
@@ -487,6 +478,13 @@ class FeaturesEnricher(TransformerMixin):
         select_features: bool, optional (default=False)
             If True, return only selected features both from input and data sources.
             Otherwise, return all features from input and only selected features from data sources.
+
+        stability_threshold: float, optional (default=0.2)
+            Stability threshold for selected features PSI calculation. If PSI is less than this threshold,
+            then feature will be dropped.
+
+        stability_agg_func: str, optional (default="max")
+            Function to aggregate stability values. Can be "max", "min", "mean".
         """
         trace_id = str(uuid.uuid4())
         if self.print_trace_id:
@@ -536,9 +534,8 @@ class FeaturesEnricher(TransformerMixin):
                     calculate_metrics=calculate_metrics,
                     estimator=estimator,
                     scoring=scoring,
-                    importance_threshold=importance_threshold,
                     stability_threshold=stability_threshold,
-                    max_features=max_features,
+                    stability_agg_func=stability_agg_func,
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
                     auto_fe_parameters=auto_fe_parameters,
                     progress_callback=progress_callback,
@@ -582,28 +579,26 @@ class FeaturesEnricher(TransformerMixin):
 
     def fit_transform(
         self,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray],
-        y: Union[pd.DataFrame, pd.Series, np.ndarray, List],
-        eval_set: Optional[Union[List[tuple], tuple]] = None,
+        X: pd.DataFrame | pd.Series | np.ndarray,
+        y: pd.DataFrame | pd.Series | np.ndarray | list,
+        eval_set: list[tuple] | tuple | None = None,
         *args,
-        exclude_features_sources: Optional[List[str]] = None,
-        keep_input: bool = True,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
-        calculate_metrics: Optional[bool] = None,
-        scoring: Union[Callable, str, None] = None,
-        estimator: Optional[Any] = None,
-        remove_outliers_calc_metrics: Optional[bool] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
+        exclude_features_sources: list[str] | None | None = None,
+        keep_input: bool | None = None,
+        calculate_metrics: bool | None = None,
+        scoring: Callable | str | None = None,
+        estimator: Any | None = None,
+        remove_outliers_calc_metrics: bool | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
         select_features: bool = True,
-        auto_fe_parameters: Optional[AutoFEParameters] = None,
-        stability_threshold: float = 0.15,
+        auto_fe_parameters: AutoFEParameters | None = None,
+        stability_threshold: float = 0.2,
+        stability_agg_func: str = "max",
         **kwargs,
     ) -> pd.DataFrame:
         """Fit to data, then transform it.
 
         Fits transformer to `X` and `y` and returns a transformed version of `X`.
-        If keep_input is True, then all input columns are copied to the output dataframe.
 
         Parameters
         ----------
@@ -613,20 +608,14 @@ class FeaturesEnricher(TransformerMixin):
         y: array-like of shape (n_samples,)
             Target values.
 
-        eval_set: List[tuple], optional (default=None)
-            List of pairs (X, y) for validation.
+        eval_set: list[tuple], optional (default=None)
+            list of pairs (X, y) for validation.
 
         keep_input: bool, optional (default=True)
+            keep_input: bool, optional (default=None)
             If True, copy original input columns to the output dataframe.
-
-        importance_threshold: float, optional (default=None)
-            Minimum SHAP value to select a feature. Default value is 0.0.
-
-        max_features: int, optional (default=None)
-            Maximum number of most important features to select. If None, the number is unlimited.
-
-        calculate_metrics: bool, optional (default=None)
-            Whether to calculate and show metrics.
+            If False, then only enriched columns are returned.
+            If None, then all search keys, ID columns, selected client features and enriched columns will be returned.
 
         estimator: sklearn-compatible estimator, optional (default=None)
             Custom estimator for metrics calculation.
@@ -642,10 +631,13 @@ class FeaturesEnricher(TransformerMixin):
             If True, return only selected features both from input and data sources.
             Otherwise, return all features from input and only selected features from data sources.
 
-        stability_threshold: float, optional (default=0.15)
+        stability_threshold: float, optional (default=0.2)
             Stability threshold for selected features PSI calculation. If PSI is less than this threshold,
             then feature will be dropped.
 
+        stability_agg_func: str, optional (default="max")
+            Function to aggregate stability values. Can be "max", "min", "mean".
+
         Returns
         -------
         X_new: pandas.DataFrame of shape (n_samples, n_features_new)
@@ -701,9 +693,8 @@ class FeaturesEnricher(TransformerMixin):
                     calculate_metrics=calculate_metrics,
                     scoring=scoring,
                     estimator=estimator,
-                    importance_threshold=importance_threshold,
                     stability_threshold=stability_threshold,
-                    max_features=max_features,
+                    stability_agg_func=stability_agg_func,
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
                     auto_fe_parameters=auto_fe_parameters,
                     progress_callback=progress_callback,
@@ -746,8 +737,6 @@ class FeaturesEnricher(TransformerMixin):
                 X,
                 exclude_features_sources=exclude_features_sources,
                 keep_input=keep_input,
-                importance_threshold=importance_threshold,
-                max_features=max_features,
                 trace_id=trace_id,
                 silent_mode=True,
                 progress_bar=progress_bar,
@@ -760,36 +749,29 @@ class FeaturesEnricher(TransformerMixin):
         self,
         X: pd.DataFrame,
         *args,
-        y: Optional[pd.Series] = None,
-        exclude_features_sources: Optional[List[str]] = None,
-        keep_input: bool = True,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
-        trace_id: Optional[str] = None,
+        y: pd.Series | None = None,
+        exclude_features_sources: list[str] | None = None,
+        keep_input: bool | None = None,
+        trace_id: str | None = None,
         metrics_calculation: bool = False,
         silent_mode=False,
-        progress_bar: Optional[ProgressBar] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
+        progress_bar: ProgressBar | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
         **kwargs,
     ) -> pd.DataFrame:
         """Transform `X`.
 
         Returns a transformed version of `X`.
-        If keep_input is True, then all input columns are copied to the output dataframe.
 
         Parameters
         ----------
         X: pandas.DataFrame of shape (n_samples, n_features)
             Input samples.
 
-        keep_input: bool, optional (default=True)
+        keep_input: bool, optional (default=None)
             If True, copy original input columns to the output dataframe.
-
-        importance_threshold: float, optional (default=None)
-            Minimum SHAP value to select a feature. Default value is 0.0.
-
-        max_features: int, optional (default=None)
-            Maximum number of most important features to select. If None, the number is unlimited.
+            If False, then only enriched columns are returned.
+            If None, then all search keys, ID columns, selected client features and enriched columns will be returned.
 
         Returns
         -------
@@ -827,11 +809,10 @@ class FeaturesEnricher(TransformerMixin):
                     X,
                     y=y,
                     exclude_features_sources=exclude_features_sources,
-                    importance_threshold=importance_threshold,
-                    max_features=max_features,
                     metrics_calculation=metrics_calculation,
                     silent_mode=silent_mode,
                     progress_bar=progress_bar,
+                    keep_input=keep_input,
                 )
                 self.logger.info("Transform finished successfully")
                 search_progress = SearchProgress(100.0, ProgressStage.FINISHED)
@@ -873,32 +854,26 @@ class FeaturesEnricher(TransformerMixin):
                     raise e
             finally:
                 self.logger.info(f"Transform elapsed time: {time.time() - start_time}")
-
-            if result is not None:
-                if keep_input:
-                    return result
-                else:
-                    return result.drop(columns=X.columns, errors="ignore")
+            
+            return result
 
     def calculate_metrics(
         self,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
-        y: Union[pd.DataFrame, pd.Series, np.ndarray, List, None] = None,
-        eval_set: Optional[Union[List[tuple], tuple]] = None,
+        X: pd.DataFrame | pd.Series | np.ndarray | None = None,
+        y: pd.DataFrame | pd.Series | np.ndarray | list | None = None,
+        eval_set: list[tuple] | tuple | None = None,
         *args,
-        scoring: Union[Callable, str, None] = None,
-        cv: Union[BaseCrossValidator, CVType, None] = None,
+        scoring: Callable | str | None = None,
+        cv: BaseCrossValidator | CVType | str | None = None,
         estimator=None,
-        exclude_features_sources: Optional[List[str]] = None,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
-        remove_outliers_calc_metrics: Optional[bool] = None,
-        trace_id: Optional[str] = None,
+        exclude_features_sources: list[str] | None = None,
+        remove_outliers_calc_metrics: bool | None = None,
+        trace_id: str | None = None,
         internal_call: bool = False,
-        progress_bar: Optional[ProgressBar] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
+        progress_bar: ProgressBar | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
         **kwargs,
-    ) -> Optional[pd.DataFrame]:
+    ) -> pd.DataFrame | None:
         """Calculate metrics
 
         Parameters
@@ -909,8 +884,8 @@ class FeaturesEnricher(TransformerMixin):
         y: array-like of shape (n_samples,), optional (default=None)
             Target values. If X not passed then y from fit will be used
 
-        eval_set: List[tuple], optional (default=None)
-            List of pairs (X, y) for validation. If X not passed then eval_set from fit will be used
+        eval_set: list[tuple], optional (default=None)
+            list of pairs (X, y) for validation. If X not passed then eval_set from fit will be used
 
         scoring: string or callable, optional (default=None)
             A string or a scorer callable object / function with signature scorer(estimator, X, y).
@@ -922,12 +897,6 @@ class FeaturesEnricher(TransformerMixin):
         estimator: sklearn-compatible estimator, optional (default=None)
             Custom estimator for metrics calculation. If not passed then CatBoost will be used.
 
-        importance_threshold: float, optional (default=None)
-            Minimum SHAP value to select a feature. Default value is 0.0.
-
-        max_features: int, optional (default=None)
-            Maximum number of most important features to select. If None, the number is unlimited.
-
         remove_outliers_calc_metrics, optional (default=True)
             If True then rows with target ouliers will be dropped on metrics calculation
 
@@ -990,8 +959,6 @@ class FeaturesEnricher(TransformerMixin):
                     validated_eval_set,
                     exclude_features_sources=exclude_features_sources,
                     cv=cv if cv is not None else self.cv,
-                    importance_threshold=importance_threshold,
-                    max_features=max_features,
                     scoring=scoring,
                     estimator=estimator,
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
@@ -1032,20 +999,19 @@ class FeaturesEnricher(TransformerMixin):
                 search_keys_for_metrics.extend([c for c in self.id_columns or [] if c not in search_keys_for_metrics])
                 self.logger.info(f"Search keys for metrics: {search_keys_for_metrics}")
 
-                prepared_data = self._prepare_data_for_metrics(
+                prepared_data = self._get_cached_enriched_data(
                     trace_id=trace_id,
                     X=X,
                     y=y,
                     eval_set=eval_set,
                     exclude_features_sources=exclude_features_sources,
-                    importance_threshold=importance_threshold,
-                    max_features=max_features,
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
                     cv_override=cv,
                     search_keys_for_metrics=search_keys_for_metrics,
                     progress_bar=progress_bar,
                     progress_callback=progress_callback,
                     client_cat_features=client_cat_features,
+                    is_for_metrics=True,
                 )
                 if prepared_data is None:
                     return None
@@ -1345,17 +1311,16 @@ class FeaturesEnricher(TransformerMixin):
     def _select_features_by_psi(
         self,
         trace_id: str,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray],
-        y: Union[pd.DataFrame, pd.Series, np.ndarray, List],
-        eval_set: Optional[Union[List[tuple], tuple]],
+        X: pd.DataFrame | pd.Series | np.ndarray,
+        y: pd.DataFrame | pd.Series | np.ndarray | list,
+        eval_set: list[tuple] | tuple | None,
         stability_threshold: float,
-        cv: Union[BaseCrossValidator, CVType, str, None] = None,
+        stability_agg_func: Callable,
+        cv: BaseCrossValidator | CVType | str | None = None,
         estimator=None,
-        exclude_features_sources: Optional[List[str]] = None,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
+        exclude_features_sources: list[str] | None = None,
         progress_bar: bool = True,
-        progress_callback: Optional[Callable] = None,
+        progress_callback: Callable | None = None,
     ):
         search_keys = self.search_keys.copy()
         validated_X, _, validated_eval_set = self._validate_train_eval(X, y, eval_set, silent=True)
@@ -1392,14 +1357,12 @@ class FeaturesEnricher(TransformerMixin):
                     c for c in client_cat_features if c not in self.id_columns_encoder.feature_names_in_
                 ]
 
-        prepared_data = self._prepare_data_for_metrics(
+        prepared_data = self._get_cached_enriched_data(
             trace_id=trace_id,
             X=X,
             y=y,
             eval_set=eval_set,
             exclude_features_sources=exclude_features_sources,
-            importance_threshold=importance_threshold,
-            max_features=max_features,
             remove_outliers_calc_metrics=False,
             cv_override=cv,
             search_keys_for_metrics=search_keys_for_metrics,
@@ -1412,15 +1375,15 @@ class FeaturesEnricher(TransformerMixin):
 
         (
             validated_X,
-            fitting_X,
+            _,
             y_sorted,
-            fitting_enriched_X,
+            _,
             _,
             fitting_eval_set_dict,
             _,
             _,
             _,
-            columns_renaming,
+            _,
             eval_set_dates,
         ) = prepared_data
 
@@ -1435,43 +1398,28 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_dates,
             search_keys,
             stability_threshold,
+            stability_agg_func,
             cat_features,
             model_task_type,
         )
-        client_features_df = self.df_with_original_index.rename(columns=columns_renaming)
-        # decoded_X = self._decode_id_columns(fitting_X, columns_renaming)
-        self._update_report_psi(trace_id, client_features_df)
 
         if unstable_features:
-            msg = f"Some features are unstable: {unstable_features} and will be dropped"
+            msg = f"{len(unstable_features)} feature(s) are unstable: {unstable_features} and will be dropped"
             self.logger.warning(msg)
             print(msg)
-            fitting_X = fitting_X.drop(columns=unstable_features, errors="ignore")
-            fitting_enriched_X = fitting_enriched_X.drop(columns=unstable_features, errors="ignore")
-            msg = f"Threre are {len(fitting_enriched_X.columns)} stable selected features left"
-            self.logger.info(msg)
-            print(msg)
-            for idx, (
-                eval_X,
-                eval_y,
-                eval_enriched_X,
-                eval_enriched_y,
-            ) in fitting_eval_set_dict.items():
-                eval_X = eval_X.drop(columns=unstable_features, errors="ignore")
-                eval_enriched_X = eval_enriched_X.drop(columns=unstable_features, errors="ignore")
-                fitting_eval_set_dict[idx] = (eval_X, eval_y, eval_enriched_X, eval_enriched_y)
 
     def _check_stability(
         self,
         X: pd.DataFrame,
-        eval_set: List[Tuple[pd.DataFrame, pd.Series]],
-        enriched_eval_set: Dict,
-        eval_set_dates: Dict[int, pd.Series],
-        search_keys: Dict[str, SearchKey],
+        eval_set: list[tuple[pd.DataFrame, pd.Series]],
+        enriched_eval_set: dict,
+        eval_set_dates: dict[int, pd.Series],
+        search_keys: dict[str, SearchKey],
         stability_threshold: float,
-        cat_features: List[str],
+        stability_agg_func: str | None,
+        cat_features: list[str],
         model_task_type: ModelTaskType,
-    ) -> List[str]:
+    ) -> list[str]:
         # Find latest eval set or earliest if all eval sets are before train set
         date_column = self._get_date_column(search_keys)
 
@@ -1521,17 +1469,17 @@ class FeaturesEnricher(TransformerMixin):
 
         unstable_by_sparsity = [feature for feature, psi in psi_values_sparse.items() if psi > stability_threshold]
         if unstable_by_sparsity:
-            self.logger.info(f"Unstable by sparsity features: {sorted(unstable_by_sparsity)}")
+            self.logger.info(f"Unstable by sparsity features ({stability_threshold}): {sorted(unstable_by_sparsity)}")
 
         psi_values = calculate_features_psi(
-            checking_eval_set_df, cat_features, date_column, self.logger, model_task_type
+            checking_eval_set_df, cat_features, date_column, self.logger, model_task_type, stability_agg_func
         )
 
         self.logger.info(f"PSI values by value: {psi_values}")
 
         unstable_by_value = [feature for feature, psi in psi_values.items() if psi > stability_threshold]
         if unstable_by_value:
-            self.logger.info(f"Unstable by value features: {sorted(unstable_by_value)}")
+            self.logger.info(f"Unstable by value features ({stability_threshold}): {sorted(unstable_by_value)}")
 
         self.psi_values = {
             feature: psi_value for feature, psi_value in psi_values.items() if psi_value <= stability_threshold
@@ -1541,7 +1489,7 @@ class FeaturesEnricher(TransformerMixin):
 
         return total_unstable_features
 
-    def _update_shap_values(self, trace_id: str, df: pd.DataFrame, new_shaps: Dict[str, float], silent: bool = False):
+    def _update_shap_values(self, trace_id: str, df: pd.DataFrame, new_shaps: dict[str, float], silent: bool = False):
         renaming = self.fit_columns_renaming or {}
         self.logger.info(f"Updating SHAP values: {new_shaps}")
         new_shaps = {
@@ -1558,7 +1506,7 @@ class FeaturesEnricher(TransformerMixin):
                 display_html_dataframe(
                     self.features_info,
                     self._features_info_without_links,
-                    self.bundle.get("relevant_features_header"),
+                    self.bundle.get("relevant_features_header").format(len(self.feature_names_)),
                     display_handle=self.features_info_display_handle,
                 )
             except (ImportError, NameError):
@@ -1596,56 +1544,6 @@ class FeaturesEnricher(TransformerMixin):
             except (ImportError, NameError):
                 pass
 
-    def _update_report_psi(self, trace_id: str, clients_features_df: pd.DataFrame):
-        self.__prepare_feature_importances(trace_id, clients_features_df)
-
-        if self.features_info_display_handle is not None:
-            try:
-                _ = get_ipython()  # type: ignore
-
-                display_html_dataframe(
-                    self.features_info,
-                    self._features_info_without_links,
-                    self.bundle.get("relevant_features_header"),
-                    display_handle=self.features_info_display_handle,
-                )
-            except (ImportError, NameError):
-                pass
-
-        if self.data_sources_display_handle is not None:
-            try:
-                _ = get_ipython()  # type: ignore
-
-                display_html_dataframe(
-                    self.relevant_data_sources,
-                    self._relevant_data_sources_wo_links,
-                    self.bundle.get("relevant_data_sources_header"),
-                    display_handle=self.data_sources_display_handle,
-                )
-            except (ImportError, NameError):
-                pass
-
-        if self.autofe_features_display_handle is not None:
-            try:
-                _ = get_ipython()  # type: ignore
-                autofe_descriptions_df = self.get_autofe_features_description()
-                if autofe_descriptions_df is not None:
-                    display_html_dataframe(
-                        df=autofe_descriptions_df,
-                        internal_df=autofe_descriptions_df,
-                        header=self.bundle.get("autofe_descriptions_header"),
-                        display_handle=self.autofe_features_display_handle,
-                    )
-            except (ImportError, NameError):
-                pass
-        if self.report_button_handle is not None:
-            try:
-                _ = get_ipython()  # type: ignore
-
-                self.__show_report_button(display_handle=self.report_button_handle)
-            except (ImportError, NameError):
-                pass
-
     def _check_train_and_eval_target_distribution(self, y, eval_set_dict):
         uneven_distribution = False
         for eval_set in eval_set_dict.values():
@@ -1659,13 +1557,13 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(msg)
 
     def _has_features_with_commercial_schema(
-        self, commercial_schema: str, exclude_features_sources: Optional[List[str]]
+        self, commercial_schema: str, exclude_features_sources: list[str] | None
     ) -> bool:
         return len(self._get_features_with_commercial_schema(commercial_schema, exclude_features_sources)) > 0
 
     def _get_features_with_commercial_schema(
-        self, commercial_schema: str, exclude_features_sources: Optional[List[str]]
-    ) -> List[str]:
+        self, commercial_schema: str, exclude_features_sources: list[str] | None
+    ) -> list[str]:
         if exclude_features_sources:
             filtered_features_info = self._internal_features_info[
                 ~self._internal_features_info[self.bundle.get("features_info_name")].isin(exclude_features_sources)
@@ -1679,15 +1577,15 @@ class FeaturesEnricher(TransformerMixin):
             ].values
         )
 
-    def _has_paid_features(self, exclude_features_sources: Optional[List[str]]) -> bool:
+    def _has_paid_features(self, exclude_features_sources: list[str] | None) -> bool:
         return self._has_features_with_commercial_schema(CommercialSchema.PAID.value, exclude_features_sources)
 
     def _is_input_same_as_fit(
         self,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
-        y: Union[pd.DataFrame, pd.Series, np.ndarray, List, None] = None,
-        eval_set: Optional[List[tuple]] = None,
-    ) -> Tuple:
+        X: pd.DataFrame | pd.Series | np.ndarray | None = None,
+        y: pd.DataFrame | pd.Series | np.ndarray | list | None = None,
+        eval_set: list[tuple] | None = None,
+    ) -> tuple:
         if X is None:
             return True, self.X, self.y, self.eval_set
 
@@ -1717,9 +1615,9 @@ class FeaturesEnricher(TransformerMixin):
     def _get_cv_and_groups(
         self,
         X: pd.DataFrame,
-        cv_override: Union[BaseCrossValidator, CVType, str, None],
-        search_keys: Dict[str, SearchKey],
-    ) -> Tuple[BaseCrossValidator, Optional[np.ndarray]]:
+        cv_override: BaseCrossValidator | CVType | str | None,
+        search_keys: dict[str, SearchKey],
+    ) -> tuple[BaseCrossValidator, np.ndarray] | None:
         _cv = cv_override or self.cv
         group_columns = sorted(self._get_group_columns(X, search_keys))
         groups = None
@@ -1747,8 +1645,8 @@ class FeaturesEnricher(TransformerMixin):
         return _cv, groups
 
     def _get_and_validate_client_cat_features(
-        self, estimator: Optional[Any], X: pd.DataFrame, search_keys: Dict[str, SearchKey]
-    ) -> Tuple[Optional[List[str]], List[str]]:
+        self, estimator: Any | None, X: pd.DataFrame, search_keys: dict[str, SearchKey]
+    ) -> tuple[list[str] | None, list[str]]:
         cat_features = []
         search_keys_for_metrics = []
         if (
@@ -1777,41 +1675,41 @@ class FeaturesEnricher(TransformerMixin):
                             raise ValidationError(self.bundle.get("cat_feature_search_key").format(cat_feature))
         return cat_features, search_keys_for_metrics
 
-    def _prepare_data_for_metrics(
+    def _get_cached_enriched_data(
         self,
         trace_id: str,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
-        y: Union[pd.DataFrame, pd.Series, np.ndarray, List, None] = None,
-        eval_set: Optional[Union[List[tuple], tuple]] = None,
-        exclude_features_sources: Optional[List[str]] = None,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
-        remove_outliers_calc_metrics: Optional[bool] = None,
-        cv_override: Union[BaseCrossValidator, CVType, str, None] = None,
-        search_keys_for_metrics: Optional[List[str]] = None,
-        progress_bar: Optional[ProgressBar] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
-        client_cat_features: Optional[List[str]] = None,
+        X: pd.DataFrame | pd.Series | np.ndarray | None = None,
+        y: pd.DataFrame | pd.Series | np.ndarray | list | None = None,
+        eval_set: list[tuple] | tuple | None = None,
+        exclude_features_sources: list[str] | None = None,
+        remove_outliers_calc_metrics: bool | None = None,
+        cv_override: BaseCrossValidator | CVType | str | None = None,
+        search_keys_for_metrics: list[str] | None = None,
+        progress_bar: ProgressBar | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
+        client_cat_features: list[str] | None = None,
+        is_for_metrics: bool = False,
     ):
         is_input_same_as_fit, X, y, eval_set = self._is_input_same_as_fit(X, y, eval_set)
         is_demo_dataset = hash_input(X, y, eval_set) in DEMO_DATASET_HASHES
         checked_eval_set = self._check_eval_set(eval_set, X)
         validated_X, validated_y, validated_eval_set = self._validate_train_eval(X, y, checked_eval_set, silent=True)
 
-        sampled_data = self._get_enriched_for_metrics(
-            trace_id,
-            validated_X,
-            validated_y,
-            validated_eval_set,
-            exclude_features_sources,
-            is_input_same_as_fit,
-            is_demo_dataset,
-            remove_outliers_calc_metrics,
-            progress_bar,
-            progress_callback,
+        sampled_data = self._get_enriched_datasets(
+            trace_id=trace_id,
+            validated_X=validated_X,
+            validated_y=validated_y,
+            eval_set=validated_eval_set,
+            exclude_features_sources=exclude_features_sources,
+            is_input_same_as_fit=is_input_same_as_fit,
+            is_demo_dataset=is_demo_dataset,
+            remove_outliers_calc_metrics=remove_outliers_calc_metrics,
+            progress_bar=progress_bar,
+            progress_callback=progress_callback,
+            is_for_metrics=is_for_metrics,
         )
-        (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming) = dataclasses.astuple(
-            sampled_data
+        (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming, generated_features) = (
+            dataclasses.astuple(sampled_data)
         )
 
         excluding_search_keys = list(search_keys.keys())
@@ -1827,14 +1725,9 @@ class FeaturesEnricher(TransformerMixin):
 
         client_features = [
             c
-            for c in X_sampled.columns.to_list()
-            if (
-                not self.fit_select_features
-                or c in set(self.feature_names_).union(self.id_columns or [])
-                or (self.fit_columns_renaming or {}).get(c, c) in set(self.feature_names_).union(self.id_columns or [])
-            )
-            and c
-            not in (
+            for c in (validated_X.columns.to_list() + generated_features)
+            if (not self.fit_select_features or c in set(self.feature_names_).union(self.id_columns or []))
+            and c not in (
                 excluding_search_keys
                 + list(self.fit_dropped_features)
                 + [DateTimeSearchKeyConverter.DATETIME_COL, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID]
@@ -1842,20 +1735,17 @@ class FeaturesEnricher(TransformerMixin):
         ]
         self.logger.info(f"Client features column on prepare data for metrics: {client_features}")
 
-        filtered_enriched_features = self.__filtered_enriched_features(
-            importance_threshold, max_features, trace_id, validated_X
-        )
-        filtered_enriched_features = [c for c in filtered_enriched_features if c not in client_features]
+        selected_enriched_features = [c for c in self.feature_names_ if c not in client_features]
 
         X_sorted, y_sorted = self._sort_by_system_record_id(X_sampled, y_sampled, self.cv)
         enriched_X_sorted, enriched_y_sorted = self._sort_by_system_record_id(enriched_X, y_sampled, self.cv)
 
         cv, groups = self._get_cv_and_groups(enriched_X_sorted, cv_override, search_keys)
 
-        existing_filtered_enriched_features = [c for c in filtered_enriched_features if c in enriched_X_sorted.columns]
+        existing_selected_enriched_features = [c for c in selected_enriched_features if c in enriched_X_sorted.columns]
 
         fitting_X = X_sorted[client_features].copy()
-        fitting_enriched_X = enriched_X_sorted[client_features + existing_filtered_enriched_features].copy()
+        fitting_enriched_X = enriched_X_sorted[client_features + existing_selected_enriched_features].copy()
 
         renamed_generate_features = [columns_renaming.get(c, c) for c in (self.generate_features or [])]
         renamed_client_cat_features = [columns_renaming.get(c, c) for c in (client_cat_features or [])]
@@ -1995,29 +1885,31 @@ class FeaturesEnricher(TransformerMixin):
         X_sampled: pd.DataFrame
         y_sampled: pd.Series
         enriched_X: pd.DataFrame
-        eval_set_sampled_dict: Dict[int, Tuple[pd.DataFrame, pd.Series]]
-        search_keys: Dict[str, SearchKey]
-        columns_renaming: Dict[str, str]
+        eval_set_sampled_dict: dict[int, tuple[pd.DataFrame, pd.Series]]
+        search_keys: dict[str, SearchKey]
+        columns_renaming: dict[str, str]
+        generated_features: list[str]
 
-    def _get_enriched_for_metrics(
+    def _get_enriched_datasets(
         self,
         trace_id: str,
-        validated_X: Union[pd.DataFrame, pd.Series, np.ndarray, None],
-        validated_y: Union[pd.DataFrame, pd.Series, np.ndarray, List, None],
-        eval_set: Optional[List[tuple]],
-        exclude_features_sources: Optional[List[str]],
+        validated_X: pd.DataFrame | pd.Series | np.ndarray | None,
+        validated_y: pd.DataFrame | pd.Series | np.ndarray | list | None,
+        eval_set: list[tuple] | None,
+        exclude_features_sources: list[str] | None,
         is_input_same_as_fit: bool,
         is_demo_dataset: bool,
-        remove_outliers_calc_metrics: Optional[bool],
-        progress_bar: Optional[ProgressBar],
-        progress_callback: Optional[Callable[[SearchProgress], Any]],
+        remove_outliers_calc_metrics: bool | None,
+        progress_bar: ProgressBar | None,
+        progress_callback: Callable[[SearchProgress], Any] | None,
+        is_for_metrics: bool = False,
     ) -> _EnrichedDataForMetrics:
         datasets_hash = hash_input(validated_X, validated_y, eval_set)
         cached_sampled_datasets = self.__cached_sampled_datasets.get(datasets_hash)
         if cached_sampled_datasets is not None and is_input_same_as_fit and remove_outliers_calc_metrics is None:
             self.logger.info("Cached enriched dataset found - use it")
             return self.__get_sampled_cached_enriched(datasets_hash, exclude_features_sources)
-        elif len(self.feature_importances_) == 0:
+        elif len(self.feature_names_) == 0 or all([f in validated_X.columns for f in self.feature_names_]):
             self.logger.info("No external features selected. So use only input datasets for metrics calculation")
             return self.__get_enriched_as_input(validated_X, validated_y, eval_set, is_demo_dataset)
         # TODO save and check if dataset was deduplicated - use imbalance branch for such case
@@ -2043,12 +1935,13 @@ class FeaturesEnricher(TransformerMixin):
                 trace_id,
                 progress_bar,
                 progress_callback,
+                is_for_metrics=is_for_metrics,
             )
 
     def __get_sampled_cached_enriched(
-        self, datasets_hash: str, exclude_features_sources: Optional[List[str]]
+        self, datasets_hash: str, exclude_features_sources: list[str] | None
     ) -> _EnrichedDataForMetrics:
-        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming = (
+        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming, generated_features = (
             self.__cached_sampled_datasets[datasets_hash]
         )
         if exclude_features_sources:
@@ -2062,10 +1955,11 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_sampled_dict,
             columns_renaming,
             search_keys,
+            generated_features,
         )
 
     def __get_enriched_as_input(
-        self, validated_X: pd.DataFrame, validated_y: pd.Series, eval_set: Optional[List[tuple]], is_demo_dataset: bool
+        self, validated_X: pd.DataFrame, validated_y: pd.Series, eval_set: list[tuple] | None, is_demo_dataset: bool
     ) -> _EnrichedDataForMetrics:
         eval_set_sampled_dict = {}
 
@@ -2130,6 +2024,10 @@ class FeaturesEnricher(TransformerMixin):
         if DateTimeSearchKeyConverter.DATETIME_COL in df.columns:
             df = df.drop(columns=DateTimeSearchKeyConverter.DATETIME_COL)
 
+        df = df.rename(columns=columns_renaming)
+        generated_features = [columns_renaming.get(c, c) for c in generated_features]
+        search_keys = {columns_renaming.get(k, k): v for k, v in search_keys.items()}
+
         train_df = df.query(f"{EVAL_SET_INDEX} == 0") if eval_set is not None else df
         X_sampled = train_df.drop(columns=[TARGET, EVAL_SET_INDEX], errors="ignore")
         y_sampled = train_df[TARGET].copy()
@@ -2152,13 +2050,14 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_sampled_dict,
             columns_renaming,
             search_keys,
+            generated_features,
         )
 
     def __get_enriched_from_fit(
         self,
-        eval_set: Optional[List[tuple]],
+        eval_set: list[tuple] | None,
         trace_id: str,
-        remove_outliers_calc_metrics: Optional[bool],
+        remove_outliers_calc_metrics: bool | None,
     ) -> _EnrichedDataForMetrics:
         eval_set_sampled_dict = {}
         search_keys = self.fit_search_keys.copy()
@@ -2246,6 +2145,7 @@ class FeaturesEnricher(TransformerMixin):
             eval_X_sampled.rename(columns=self.fit_columns_renaming, inplace=True)
             enriched_eval_X.rename(columns=self.fit_columns_renaming, inplace=True)
         search_keys = {self.fit_columns_renaming.get(k, k): v for k, v in search_keys.items()}
+        generated_features = [self.fit_columns_renaming.get(c, c) for c in self.fit_generated_features]
 
         datasets_hash = hash_input(self.X, self.y, self.eval_set)
         return self.__cache_and_return_results(
@@ -2256,17 +2156,19 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_sampled_dict,
             self.fit_columns_renaming,
             search_keys,
+            generated_features,
         )
 
     def __get_enriched_from_transform(
         self,
         validated_X: pd.DataFrame,
         validated_y: pd.Series,
-        eval_set: Optional[List[tuple]],
-        exclude_features_sources: Optional[List[str]],
+        eval_set: list[tuple] | None,
+        exclude_features_sources: list[str] | None,
         trace_id: str,
-        progress_bar: Optional[ProgressBar],
-        progress_callback: Optional[Callable[[SearchProgress], Any]],
+        progress_bar: ProgressBar | None,
+        progress_callback: Callable[[SearchProgress], Any] | None,
+        is_for_metrics: bool = False,
     ) -> _EnrichedDataForMetrics:
         has_eval_set = eval_set is not None
 
@@ -2274,6 +2176,16 @@ class FeaturesEnricher(TransformerMixin):
 
         # Prepare
         df = self.__combine_train_and_eval_sets(validated_X, validated_y, eval_set)
+
+        # Exclude OOT eval sets from transform because they are not used for metrics calculation
+        if not is_for_metrics and EVAL_SET_INDEX in df.columns:
+            for eval_index in df[EVAL_SET_INDEX].unique():
+                if eval_index == 0:
+                    continue
+                eval_df = df.query(f"{EVAL_SET_INDEX} == {eval_index}")
+                if eval_df[TARGET].isna().all():
+                    df = df.query(f"{EVAL_SET_INDEX} != {eval_index}")
+
         df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
         df = self.__downsample_for_metrics(df)
 
@@ -2315,10 +2227,11 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_sampled_dict,
             columns_renaming,
             search_keys,
+            generated_features,
         )
 
     def __combine_train_and_eval_sets(
-        self, X: pd.DataFrame, y: Optional[pd.Series] = None, eval_set: Optional[List[tuple]] = None
+        self, X: pd.DataFrame, y: pd.Series | None = None, eval_set: list[tuple] | None = None
     ) -> pd.DataFrame:
         df = X.copy()
         if y is not None:
@@ -2370,8 +2283,8 @@ class FeaturesEnricher(TransformerMixin):
         )
 
     def __extract_train_data(
-        self, enriched_df: pd.DataFrame, x_columns: List[str]
-    ) -> Tuple[pd.DataFrame, pd.Series, pd.DataFrame]:
+        self, enriched_df: pd.DataFrame, x_columns: list[str]
+    ) -> tuple[pd.DataFrame, pd.Series, pd.DataFrame]:
         if EVAL_SET_INDEX in enriched_df.columns:
             enriched_Xy = enriched_df.query(f"{EVAL_SET_INDEX} == 0")
         else:
@@ -2382,8 +2295,8 @@ class FeaturesEnricher(TransformerMixin):
         return X_sampled, y_sampled, enriched_X
 
     def __extract_eval_data(
-        self, enriched_df: pd.DataFrame, x_columns: List[str], enriched_X_columns: List[str], eval_set_len: int
-    ) -> Tuple[Dict[int, Tuple], Dict[int, pd.Series]]:
+        self, enriched_df: pd.DataFrame, x_columns: list[str], enriched_X_columns: list[str], eval_set_len: int
+    ) -> tuple[dict[int, tuple], dict[int, pd.Series]]:
         eval_set_sampled_dict = {}
 
         for idx in range(eval_set_len):
@@ -2401,9 +2314,10 @@ class FeaturesEnricher(TransformerMixin):
         X_sampled: pd.DataFrame,
         y_sampled: pd.Series,
         enriched_X: pd.DataFrame,
-        eval_set_sampled_dict: Dict[int, Tuple],
-        columns_renaming: Dict[str, str],
-        search_keys: Dict[str, SearchKey],
+        eval_set_sampled_dict: dict[int, tuple],
+        columns_renaming: dict[str, str],
+        search_keys: dict[str, SearchKey],
+        generated_features: list[str],
     ) -> _EnrichedDataForMetrics:
 
         self.__cached_sampled_datasets[datasets_hash] = (
@@ -2413,10 +2327,11 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_sampled_dict,
             search_keys,
             columns_renaming,
+            generated_features,
         )
 
         return self.__mk_sampled_data_tuple(
-            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming
+            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming, generated_features
         )
 
     def __mk_sampled_data_tuple(
@@ -2424,17 +2339,11 @@ class FeaturesEnricher(TransformerMixin):
         X_sampled: pd.DataFrame,
         y_sampled: pd.Series,
         enriched_X: pd.DataFrame,
-        eval_set_sampled_dict: Dict,
-        search_keys: Dict,
-        columns_renaming: Dict[str, str],
+        eval_set_sampled_dict: dict,
+        search_keys: dict,
+        columns_renaming: dict[str, str],
+        generated_features: list[str],
     ):
-        # X_sampled - with hash-suffixes
-        # reversed_renaming = {v: k for k, v in columns_renaming.items()}
-        # search_keys = {
-        #     reversed_renaming.get(k, k): v
-        #     for k, v in search_keys.items()
-        #     if reversed_renaming.get(k, k) in X_sampled.columns.to_list()
-        # }
         return FeaturesEnricher._EnrichedDataForMetrics(
             X_sampled=X_sampled,
             y_sampled=y_sampled,
@@ -2442,9 +2351,10 @@ class FeaturesEnricher(TransformerMixin):
             eval_set_sampled_dict=eval_set_sampled_dict,
             search_keys=search_keys,
             columns_renaming=columns_renaming,
+            generated_features=generated_features,
         )
 
-    def get_search_id(self) -> Optional[str]:
+    def get_search_id(self) -> str | None:
         """Returns search_id of the fitted enricher. Not available before a successful fit."""
         return self._search_task.search_task_id if self._search_task else None
 
@@ -2457,7 +2367,7 @@ class FeaturesEnricher(TransformerMixin):
 
         return self.features_info
 
-    def get_progress(self, trace_id: Optional[str] = None, search_task: Optional[SearchTask] = None) -> SearchProgress:
+    def get_progress(self, trace_id: str | None = None, search_task: SearchTask | None = None) -> SearchProgress:
         search_task = search_task or self._search_task
         if search_task is not None:
             trace_id = trace_id or uuid.uuid4()
@@ -2565,16 +2475,15 @@ if response.status_code == 200:
         trace_id: str,
         X: pd.DataFrame,
         *,
-        y: Optional[pd.Series] = None,
-        exclude_features_sources: Optional[List[str]] = None,
-        importance_threshold: Optional[float] = None,
-        max_features: Optional[int] = None,
+        y: pd.Series | None = None,
+        exclude_features_sources: list[str] | None = None,
         metrics_calculation: bool = False,
         silent_mode: bool = False,
-        progress_bar: Optional[ProgressBar] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
+        progress_bar: ProgressBar | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
         add_fit_system_record_id: bool = False,
-    ) -> Tuple[pd.DataFrame, Dict[str, str], List[str], Dict[str, SearchKey]]:
+        keep_input: bool | None = None,
+    ) -> tuple[pd.DataFrame, dict[str, str], list[str], dict[str, SearchKey]]:
         if self._search_task is None:
             raise NotFittedError(self.bundle.get("transform_unfitted_enricher"))
 
@@ -2592,11 +2501,8 @@ if response.status_code == 200:
 
             self.__log_debug_information(validated_X, validated_y, exclude_features_sources=exclude_features_sources)
 
-            filtered_columns = self.__filtered_enriched_features(
-                importance_threshold, max_features, trace_id, validated_X
-            )
             # If there are no important features, return original dataframe
-            if not filtered_columns:
+            if len(self.feature_names_) == 0:
                 msg = self.bundle.get("no_important_features_for_transform")
                 self.__log_warning(msg, show_support_link=True)
                 return X, {c: c for c in X.columns}, [], dict()
@@ -2703,9 +2609,9 @@ if response.status_code == 200:
             if not external_features:
                 self.logger.warning(
                     "No external features found, returning original dataframe"
-                    f" with generated important features: {filtered_columns}"
+                    f" with generated important features: {self.feature_names_}"
                 )
-                filtered_columns = [c for c in filtered_columns if c in df.columns]
+                filtered_columns = [c for c in self.feature_names_ if c in df.columns]
                 self.logger.warning(f"Filtered columns by existance in dataframe: {filtered_columns}")
                 return df[filtered_columns], columns_renaming, generated_features, search_keys
 
@@ -2843,16 +2749,6 @@ if response.status_code == 200:
             )
             dataset.columns_renaming = columns_renaming
 
-            if max_features is not None or importance_threshold is not None:
-                exclude_features_sources = list(
-                    set(
-                        (exclude_features_sources or [])
-                        + self._get_excluded_features(max_features, importance_threshold)
-                    )
-                )
-                if len(exclude_features_sources) == 0:
-                    exclude_features_sources = None
-
             validation_task = self._search_task.validation(
                 trace_id,
                 dataset,
@@ -2880,7 +2776,7 @@ if response.status_code == 200:
                 progress_bar.progress = progress.to_progress_bar()
             if progress_callback is not None:
                 progress_callback(progress)
-            prev_progress: Optional[SearchProgress] = None
+            prev_progress: SearchProgress | None = None
             polling_period_seconds = 1
             try:
                 while progress.stage != ProgressStage.DOWNLOADING.value:
@@ -2917,6 +2813,8 @@ if response.status_code == 200:
                 print(self.bundle.get("transform_start"))
 
             # Prepare input DataFrame for __enrich by concatenating generated ids and client features
+            df_before_explode = df_before_explode.rename(columns=columns_renaming)
+            generated_features = [columns_renaming.get(c, c) for c in generated_features]
             combined_df = pd.concat(
                 [
                     validated_Xy.reset_index(drop=True),
@@ -2934,15 +2832,28 @@ if response.status_code == 200:
             )
 
             selected_generated_features = [
-                c for c in generated_features if not self.fit_select_features or c in filtered_columns
-            ]
-            selecting_columns = [
-                c
-                for c in itertools.chain(validated_Xy.columns.tolist(), selected_generated_features)
-                if (c not in self.zero_shap_client_features and c not in self.unstable_client_features)
-                or c in (self.id_columns or [])
+                c for c in generated_features if not self.fit_select_features or c in self.feature_names_
             ]
-            selecting_columns.extend(c for c in result.columns if c in filtered_columns and c not in selecting_columns)
+            if keep_input is None:
+                selected_input_columns = [
+                    c
+                    for c in validated_Xy.columns
+                    if not self.fit_select_features
+                    or c in self.feature_names_
+                    or c in self.search_keys
+                    or c in (self.id_columns or [])
+                    or c in [EVAL_SET_INDEX, TARGET]  # transform for metrics calculation
+                ]
+            elif keep_input is True:
+                selected_input_columns = validated_Xy.columns.to_list()
+            else:
+                selected_input_columns = []
+
+            selecting_columns = selected_input_columns + selected_generated_features
+            selecting_columns.extend(
+                c for c in result.columns
+                if c in self.feature_names_ and c not in selecting_columns and c not in validated_Xy.columns
+            )
             if add_fit_system_record_id:
                 selecting_columns.append(SORT_ID)
 
@@ -2968,29 +2879,7 @@ if response.status_code == 200:
 
             return result, columns_renaming, generated_features, search_keys
 
-    def _get_excluded_features(self, max_features: Optional[int], importance_threshold: Optional[float]) -> List[str]:
-        features_info = self._internal_features_info
-        comm_schema_header = self.bundle.get("features_info_commercial_schema")
-        shap_value_header = self.bundle.get("features_info_shap")
-        feature_name_header = self.bundle.get("features_info_name")
-        external_features = features_info[features_info[comm_schema_header].str.len() > 0]
-        filtered_features = external_features
-        if importance_threshold is not None:
-            filtered_features = filtered_features[filtered_features[shap_value_header] >= importance_threshold]
-        if max_features is not None and len(filtered_features) > max_features:
-            filtered_features = filtered_features.iloc[:max_features, :]
-        if len(filtered_features) == len(external_features):
-            return []
-        else:
-            if len(filtered_features[filtered_features[comm_schema_header].isin([CommercialSchema.PAID.value])]):
-                return []
-            excluded_features = external_features[~external_features.index.isin(filtered_features.index)].copy()
-            excluded_features = excluded_features[
-                excluded_features[comm_schema_header].isin([CommercialSchema.PAID.value])
-            ]
-            return excluded_features[feature_name_header].values.tolist()
-
-    def __validate_search_keys(self, search_keys: Dict[str, SearchKey], search_id: Optional[str] = None):
+    def __validate_search_keys(self, search_keys: dict[str, SearchKey], search_id: str | None = None):
         if (search_keys is None or len(search_keys) == 0) and self.country_code is None:
             if search_id:
                 self.logger.debug(f"search_id {search_id} provided without search_keys")
@@ -3063,23 +2952,22 @@ if response.status_code == 200:
     def __inner_fit(
         self,
         trace_id: str,
-        X: Union[pd.DataFrame, pd.Series, np.ndarray],
-        y: Union[pd.DataFrame, pd.Series, np.ndarray, List, None],
-        eval_set: Optional[List[tuple]],
-        progress_bar: Optional[ProgressBar],
+        X: pd.DataFrame | pd.Series | np.ndarray,
+        y: pd.DataFrame | pd.Series | np.ndarray | list | None,
+        eval_set: list[tuple] | None,
+        progress_bar: ProgressBar | None,
         start_time: int,
         *,
-        exclude_features_sources: Optional[List[str]] = None,
-        calculate_metrics: Optional[bool],
-        scoring: Union[Callable, str, None],
-        estimator: Optional[Any],
-        importance_threshold: Optional[float],
+        exclude_features_sources: list[str] | None = None,
+        calculate_metrics: bool | None,
+        scoring: Callable | str | None,
+        estimator: Any | None,
         stability_threshold: float,
-        max_features: Optional[int],
-        remove_outliers_calc_metrics: Optional[bool],
+        stability_agg_func: str,
+        remove_outliers_calc_metrics: bool | None,
         auto_fe_parameters: AutoFEParameters,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
-        search_id_callback: Optional[Callable[[str], Any]] = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
+        search_id_callback: Callable[[str], Any] | None = None,
     ):
         self._search_task = None
         self.warning_counter.reset()
@@ -3140,7 +3028,6 @@ if response.status_code == 200:
         )
 
         df = self.__combine_train_and_eval_sets(validated_X, validated_y, validated_eval_set)
-        self.id_columns_encoder = OrdinalEncoder().fit(df[self.id_columns or []])
 
         self.fit_search_keys = self.search_keys.copy()
         df = self.__handle_index_search_keys(df, self.fit_search_keys)
@@ -3148,8 +3035,22 @@ if response.status_code == 200:
 
         maybe_date_column = SearchKey.find_key(self.fit_search_keys, [SearchKey.DATE, SearchKey.DATETIME])
         has_date = maybe_date_column is not None
+
         self.model_task_type = self.model_task_type or define_task(validated_y, has_date, self.logger)
 
+        if EVAL_SET_INDEX in df.columns:
+            only_train_df = df.query(f"{EVAL_SET_INDEX} == 0")
+            only_train_df = only_train_df.drop(columns=[EVAL_SET_INDEX])
+        else:
+            only_train_df = df
+
+        self.imbalanced = is_imbalanced(only_train_df, self.model_task_type, self.sample_config, self.bundle)
+        if self.imbalanced:
+            # Exclude eval sets from fit because they will be transformed before metrics calculation
+            df = only_train_df
+
+        self.id_columns_encoder = OrdinalEncoder().fit(df[self.id_columns or []])
+
         self._validate_binary_observations(validated_y, self.model_task_type)
 
         self.runtime_parameters = get_runtime_params_custom_loss(
@@ -3365,6 +3266,7 @@ if response.status_code == 200:
             model_task_type=self.model_task_type,
             cv_type=self.cv,
             id_columns=self.__get_renamed_id_columns(),
+            is_imbalanced=self.imbalanced,
             date_column=self._get_date_column(self.fit_search_keys),
             date_format=self.date_format,
             random_state=self.random_state,
@@ -3444,8 +3346,6 @@ if response.status_code == 200:
             if progress_callback is not None:
                 progress_callback(progress)
 
-            self.imbalanced = dataset.imbalanced
-
             zero_hit_search_keys = self._search_task.get_zero_hit_rate_search_keys()
             if zero_hit_search_keys:
                 self.logger.warning(
@@ -3468,33 +3368,24 @@ if response.status_code == 200:
 
             self.__prepare_feature_importances(trace_id, df)
 
-            self.__show_selected_features(self.fit_search_keys)
-
-            autofe_description = self.get_autofe_features_description()
-            if autofe_description is not None and len(autofe_description) > 0:
-                self.logger.info(f"AutoFE descriptions: {autofe_description}")
-                self.autofe_features_display_handle = display_html_dataframe(
-                    df=autofe_description,
-                    internal_df=autofe_description,
-                    header=self.bundle.get("autofe_descriptions_header"),
-                    display_id=f"autofe_descriptions_{uuid.uuid4()}",
-                )
-
             self._select_features_by_psi(
                 trace_id=trace_id,
                 X=X,
                 y=y,
                 eval_set=eval_set,
                 stability_threshold=stability_threshold,
+                stability_agg_func=stability_agg_func,
                 cv=self.cv,
                 estimator=estimator,
                 exclude_features_sources=exclude_features_sources,
-                importance_threshold=importance_threshold,
-                max_features=max_features,
                 progress_bar=progress_bar,
                 progress_callback=progress_callback,
             )
 
+            self.__prepare_feature_importances(trace_id, df)
+
+            self.__show_selected_features()
+
             if self._has_paid_features(exclude_features_sources):
                 if calculate_metrics is not None and calculate_metrics:
                     msg = self.bundle.get("metrics_with_paid_features")
@@ -3525,8 +3416,6 @@ if response.status_code == 200:
                         self.__show_metrics(
                             scoring,
                             estimator,
-                            importance_threshold,
-                            max_features,
                             remove_outliers_calc_metrics,
                             trace_id,
                             progress_bar,
@@ -3543,7 +3432,7 @@ if response.status_code == 200:
             if not self.warning_counter.has_warnings():
                 self.__display_support_link(self.bundle.get("all_ok_community_invite"))
 
-    def __convert_unnestable_keys(self, df: pd.DataFrame, unnest_search_keys: Dict[str, str]):
+    def __convert_unnestable_keys(self, df: pd.DataFrame, unnest_search_keys: dict[str, str]):
         email_column = self._get_email_column(self.fit_search_keys)
         hem_column = self._get_hem_column(self.fit_search_keys)
         if email_column:
@@ -3575,7 +3464,7 @@ if response.status_code == 200:
     def __should_add_date_column(self):
         return self.add_date_if_missing or (self.cv is not None and self.cv.is_time_series())
 
-    def __get_renamed_id_columns(self, renaming: Optional[Dict[str, str]] = None):
+    def __get_renamed_id_columns(self, renaming: dict[str, str] | None = None):
         renaming = renaming or self.fit_columns_renaming
         reverse_renaming = {v: k for k, v in renaming.items()}
         return None if self.id_columns is None else [reverse_renaming.get(c) or c for c in self.id_columns]
@@ -3609,7 +3498,7 @@ if response.status_code == 200:
         self.cv = cv
         self.runtime_parameters.properties["cv_type"] = self.cv.name
 
-    def get_columns_by_search_keys(self, keys: List[str]):
+    def get_columns_by_search_keys(self, keys: list[str]):
         if "HEM" in keys:
             keys.append("EMAIL")
         if "DATE" in keys:
@@ -3620,11 +3509,11 @@ if response.status_code == 200:
     def _validate_train_eval(
         self,
         X: pd.DataFrame,
-        y: Optional[pd.Series] = None,
-        eval_set: Optional[List[Tuple[pd.DataFrame, pd.Series]]] = None,
+        y: pd.Series | None = None,
+        eval_set: list[tuple[pd.DataFrame, pd.Series]] | None = None,
         is_transform: bool = False,
         silent: bool = False,
-    ) -> Tuple[pd.DataFrame, pd.Series, Optional[List[Tuple[pd.DataFrame, pd.Series]]]]:
+    ) -> tuple[pd.DataFrame, pd.Series, list[tuple[pd.DataFrame, pd.Series]]] | None:
         validated_X = self._validate_X(X, is_transform)
         validated_y = self._validate_y(validated_X, y, enforce_y=not is_transform)
         validated_eval_set = self._validate_eval_set(validated_X, eval_set, silent=silent)
@@ -3633,7 +3522,7 @@ if response.status_code == 200:
     def _encode_id_columns(
         self,
         X: pd.DataFrame,
-    ) -> Tuple[pd.DataFrame, Dict[str, List[Any]]]:
+    ) -> tuple[pd.DataFrame, dict[str, list[Any]]]:
         unknown_dict = {}
 
         if self.id_columns and self.id_columns_encoder is not None:
@@ -3703,7 +3592,7 @@ if response.status_code == 200:
 
         return validated_X
 
-    def _validate_y(self, X: pd.DataFrame, y, enforce_y: bool = True) -> Optional[pd.Series]:
+    def _validate_y(self, X: pd.DataFrame, y, enforce_y: bool = True) -> pd.Series | None:
         if y is None and not enforce_y:
             return None
         if (
@@ -3753,7 +3642,7 @@ if response.status_code == 200:
         return validated_y
 
     def _validate_eval_set(
-        self, X: pd.DataFrame, eval_set: Optional[List[Tuple[pd.DataFrame, pd.Series]]], silent: bool = False
+        self, X: pd.DataFrame, eval_set: list[tuple[pd.DataFrame, pd.Series]] | None, silent: bool = False
     ):
         if eval_set is None:
             return None
@@ -3777,7 +3666,7 @@ if response.status_code == 200:
 
         return validated_eval_set
 
-    def _validate_eval_set_pair(self, X: pd.DataFrame, eval_pair: Tuple) -> Tuple[pd.DataFrame, pd.Series]:
+    def _validate_eval_set_pair(self, X: pd.DataFrame, eval_pair: tuple) -> tuple[pd.DataFrame, pd.Series]:
         if len(eval_pair) != 2:
             raise ValidationError(self.bundle.get("eval_set_invalid_tuple_size").format(len(eval_pair)))
         eval_X, eval_y = eval_pair
@@ -3865,7 +3754,7 @@ if response.status_code == 200:
 
         return validated_eval_X, validated_eval_y
 
-    def _validate_baseline_score(self, X: pd.DataFrame, eval_set: Optional[List[Tuple]]):
+    def _validate_baseline_score(self, X: pd.DataFrame, eval_set: list[tuple] | None):
         if self.baseline_score_column is not None:
             if self.baseline_score_column not in X.columns:
                 raise ValidationError(
@@ -3885,15 +3774,15 @@ if response.status_code == 200:
                             raise ValidationError(self.bundle.get("baseline_score_column_has_na"))
 
     @staticmethod
-    def _sample_X_and_y(X: pd.DataFrame, y: pd.Series, enriched_X: pd.DataFrame) -> Tuple[pd.DataFrame, pd.Series]:
+    def _sample_X_and_y(X: pd.DataFrame, y: pd.Series, enriched_X: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series]:
         Xy = pd.concat([X, y], axis=1)
         Xy = pd.merge(Xy, enriched_X, left_index=True, right_index=True, how="inner", suffixes=("", "enriched"))
         return Xy[X.columns].copy(), Xy[TARGET].copy()
 
     @staticmethod
     def _sort_by_system_record_id(
-        X: pd.DataFrame, y: pd.Series, cv: Optional[CVType]
-    ) -> Tuple[pd.DataFrame, pd.Series]:
+        X: pd.DataFrame, y: pd.Series, cv: CVType | None
+    ) -> tuple[pd.DataFrame, pd.Series]:
         if cv not in [CVType.time_series, CVType.blocked_time_series]:
             record_id_column = ENTITY_SYSTEM_RECORD_ID if ENTITY_SYSTEM_RECORD_ID in X else SYSTEM_RECORD_ID
             Xy = X.copy()
@@ -3910,8 +3799,8 @@ if response.status_code == 200:
     # Deprecated
     @staticmethod
     def _sort_by_keys(
-        X: pd.DataFrame, y: pd.Series, search_keys: Dict[str, SearchKey], cv: Optional[CVType]
-    ) -> Tuple[pd.DataFrame, pd.Series]:
+        X: pd.DataFrame, y: pd.Series, search_keys: dict[str, SearchKey], cv: CVType | None
+    ) -> tuple[pd.DataFrame, pd.Series]:
         if cv not in [CVType.time_series, CVType.blocked_time_series]:
             if DateTimeSearchKeyConverter.DATETIME_COL in X.columns:
                 date_column = DateTimeSearchKeyConverter.DATETIME_COL
@@ -3950,16 +3839,14 @@ if response.status_code == 200:
     def __log_debug_information(
         self,
         X: pd.DataFrame,
-        y: Union[pd.Series, np.ndarray, list, None] = None,
-        eval_set: Optional[List[tuple]] = None,
-        exclude_features_sources: Optional[List[str]] = None,
-        calculate_metrics: Optional[bool] = None,
-        cv: Optional[Any] = None,
-        importance_threshold: Optional[Any] = None,
-        max_features: Optional[Any] = None,
-        scoring: Optional[Any] = None,
-        estimator: Optional[Any] = None,
-        remove_outliers_calc_metrics: Optional[bool] = None,
+        y: pd.Series | np.ndarray | list | None = None,
+        eval_set: list[tuple] | None = None,
+        exclude_features_sources: list[str] | None = None,
+        calculate_metrics: bool | None = None,
+        cv: Any | None = None,
+        scoring: Any | None = None,
+        estimator: Any | None = None,
+        remove_outliers_calc_metrics: bool | None = None,
     ):
         try:
             resolved_api_key = self.api_key or os.environ.get(UPGINI_API_KEY)
@@ -3972,8 +3859,6 @@ if response.status_code == 200:
                 f"Runtime parameters: {self.runtime_parameters}\n"
                 f"Date format: {self.date_format}\n"
                 f"CV: {cv}\n"
-                f"importance_threshold: {importance_threshold}\n"
-                f"max_features: {max_features}\n"
                 f"Shared datasets: {self.shared_datasets}\n"
                 f"Random state: {self.random_state}\n"
                 f"Generate features: {self.generate_features}\n"
@@ -4037,7 +3922,7 @@ if response.status_code == 200:
         except Exception:
             self.logger.warning("Failed to log debug information", exc_info=True)
 
-    def __handle_index_search_keys(self, df: pd.DataFrame, search_keys: Dict[str, SearchKey]) -> pd.DataFrame:
+    def __handle_index_search_keys(self, df: pd.DataFrame, search_keys: dict[str, SearchKey]) -> pd.DataFrame:
         index_names = df.index.names if df.index.names != [None] else [DEFAULT_INDEX]
         index_search_keys = set(index_names).intersection(search_keys.keys())
         if len(index_search_keys) > 0:
@@ -4056,7 +3941,7 @@ if response.status_code == 200:
         return df
 
     def _add_current_date_as_key(
-        self, df: pd.DataFrame, search_keys: Dict[str, SearchKey], bundle: ResourceBundle, silent: bool = False
+        self, df: pd.DataFrame, search_keys: dict[str, SearchKey], bundle: ResourceBundle, silent: bool = False
     ) -> pd.DataFrame:
         if (
             set(search_keys.values()) == {SearchKey.PHONE}
@@ -4073,7 +3958,7 @@ if response.status_code == 200:
         return df
 
     @staticmethod
-    def _get_group_columns(df: pd.DataFrame, search_keys: Dict[str, SearchKey]) -> List[str]:
+    def _get_group_columns(df: pd.DataFrame, search_keys: dict[str, SearchKey]) -> list[str]:
         search_key_priority = [SearchKey.PHONE, SearchKey.EMAIL, SearchKey.HEM, SearchKey.IP]
         for key_type in search_key_priority:
             if key_type in search_keys.values():
@@ -4086,7 +3971,7 @@ if response.status_code == 200:
         ]
 
     @staticmethod
-    def _get_email_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_email_column(search_keys: dict[str, SearchKey]) -> str | None:
         cols = [col for col, t in search_keys.items() if t == SearchKey.EMAIL]
         if len(cols) > 1:
             raise Exception("More than one email column found after unnest")
@@ -4094,7 +3979,7 @@ if response.status_code == 200:
             return cols[0]
 
     @staticmethod
-    def _get_hem_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_hem_column(search_keys: dict[str, SearchKey]) -> str | None:
         cols = [col for col, t in search_keys.items() if t == SearchKey.HEM]
         if len(cols) > 1:
             raise Exception("More than one hem column found after unnest")
@@ -4102,7 +3987,7 @@ if response.status_code == 200:
             return cols[0]
 
     @staticmethod
-    def _get_ip_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_ip_column(search_keys: dict[str, SearchKey]) -> str | None:
         cols = [col for col, t in search_keys.items() if t == SearchKey.IP]
         if len(cols) > 1:
             raise Exception("More than one ip column found after unnest")
@@ -4110,32 +3995,32 @@ if response.status_code == 200:
             return cols[0]
 
     @staticmethod
-    def _get_phone_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_phone_column(search_keys: dict[str, SearchKey]) -> str | None:
         for col, t in search_keys.items():
             if t == SearchKey.PHONE:
                 return col
 
     @staticmethod
-    def _get_country_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_country_column(search_keys: dict[str, SearchKey]) -> str | None:
         for col, t in search_keys.items():
             if t == SearchKey.COUNTRY:
                 return col
 
     @staticmethod
-    def _get_postal_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_postal_column(search_keys: dict[str, SearchKey]) -> str | None:
         for col, t in search_keys.items():
             if t == SearchKey.POSTAL_CODE:
                 return col
 
     @staticmethod
-    def _get_date_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
+    def _get_date_column(search_keys: dict[str, SearchKey]) -> str | None:
         return SearchKey.find_key(search_keys, [SearchKey.DATE, SearchKey.DATETIME])
 
     def _explode_multiple_search_keys(
-        self, df: pd.DataFrame, search_keys: Dict[str, SearchKey], columns_renaming: Dict[str, str]
-    ) -> Tuple[pd.DataFrame, Dict[str, List[str]]]:
+        self, df: pd.DataFrame, search_keys: dict[str, SearchKey], columns_renaming: dict[str, str]
+    ) -> tuple[pd.DataFrame, dict[str, list[str]]]:
         # find groups of multiple search keys
-        search_key_names_by_type: Dict[SearchKey, List[str]] = {}
+        search_key_names_by_type: dict[SearchKey, list[str]] = {}
         for key_name, key_type in search_keys.items():
             search_key_names_by_type[key_type] = search_key_names_by_type.get(key_type, []) + [key_name]
         search_key_names_by_type = {
@@ -4171,14 +4056,14 @@ if response.status_code == 200:
     @staticmethod
     def _add_fit_system_record_id(
         df: pd.DataFrame,
-        search_keys: Dict[str, SearchKey],
+        search_keys: dict[str, SearchKey],
         id_name: str,
         target_name: str,
-        columns_renaming: Dict[str, str],
-        id_columns: Optional[List[str]],
-        cv: Optional[CVType],
+        columns_renaming: dict[str, str],
+        id_columns: list[str] | None,
+        cv: CVType | None,
         model_task_type: ModelTaskType,
-        logger: Optional[logging.Logger] = None,
+        logger: logging.Logger | None = None,
         bundle: ResourceBundle = bundle,
     ) -> pd.DataFrame:
         original_index_name = df.index.name
@@ -4296,7 +4181,7 @@ if response.status_code == 200:
 
         return df
 
-    def __add_country_code(self, df: pd.DataFrame, search_keys: Dict[str, SearchKey]) -> pd.DataFrame:
+    def __add_country_code(self, df: pd.DataFrame, search_keys: dict[str, SearchKey]) -> pd.DataFrame:
         self.country_added = False
 
         if self.country_code is not None and SearchKey.COUNTRY not in search_keys.values():
@@ -4314,7 +4199,7 @@ if response.status_code == 200:
     def __enrich(
         self,
         input_df: pd.DataFrame,
-        result_features: Optional[pd.DataFrame],
+        result_features: pd.DataFrame | None,
         how: str = "inner",
         drop_system_record_id=True,
     ) -> pd.DataFrame:
@@ -4422,7 +4307,7 @@ if response.status_code == 200:
 
         return importances
 
-    def __get_categorical_features(self) -> List[str]:
+    def __get_categorical_features(self) -> list[str]:
         features_meta = self._search_task.get_all_features_metadata_v2()
         if features_meta is None:
             raise Exception(self.bundle.get("missing_features_meta"))
@@ -4433,11 +4318,13 @@ if response.status_code == 200:
         self,
         trace_id: str,
         clients_features_df: pd.DataFrame,
-        updated_shaps: Optional[Dict[str, float]] = None,
+        updated_shaps: dict[str, float] | None = None,
+        update_selected_features: bool = True,
         silent=False,
     ):
         if self._search_task is None:
             raise NotFittedError(self.bundle.get("transform_unfitted_enricher"))
+        selected_features = self._search_task.get_selected_features(trace_id)
         features_meta = self._search_task.get_all_features_metadata_v2()
         if features_meta is None:
             raise Exception(self.bundle.get("missing_features_meta"))
@@ -4451,8 +4338,6 @@ if response.status_code == 200:
 
         self.feature_names_ = []
         self.external_source_feature_names = []
-        self.zero_shap_client_features = []
-        self.unstable_client_features = []
         self.feature_importances_ = []
         features_info = []
         features_info_without_links = []
@@ -4460,12 +4345,19 @@ if response.status_code == 200:
 
         original_shaps = {original_names_dict.get(fm.name, fm.name): fm.shap_value for fm in features_meta}
 
+        selected_features_meta = []
         for feature_meta in features_meta:
             original_name = original_names_dict.get(feature_meta.name, feature_meta.name)
             feature_meta.name = original_name
 
             is_client_feature = original_name in clients_features_df.columns
 
+            if selected_features is not None and feature_meta.name not in selected_features:
+                self.logger.info(f"Feature {feature_meta.name} is not selected before and skipped")
+                continue
+
+            selected_features_meta.append(feature_meta)
+
             # Show and update shap values for client features only if select_features is True
             if updated_shaps is not None and (not is_client_feature or self.fit_select_features):
                 updating_shap = updated_shaps.get(feature_meta.name)
@@ -4477,9 +4369,9 @@ if response.status_code == 200:
                     updating_shap = 0.0
                 feature_meta.shap_value = updating_shap
 
-        features_meta.sort(key=lambda m: (-m.shap_value, m.name))
+        selected_features_meta.sort(key=lambda m: (-m.shap_value, m.name))
 
-        for feature_meta in features_meta:
+        for feature_meta in selected_features_meta:
             original_name = original_names_dict.get(feature_meta.name, feature_meta.name)
             is_client_feature = original_name in clients_features_df.columns
 
@@ -4490,15 +4382,11 @@ if response.status_code == 200:
                 if original_name in self.psi_values:
                     feature_meta.psi_value = self.psi_values[original_name]
                 else:
-                    if is_client_feature and self.fit_select_features:
-                        self.unstable_client_features.append(original_name)
                     continue
 
             # TODO make a decision about selected features based on special flag from mlb
 
             if original_shaps.get(feature_meta.name, 0.0) == 0.0:
-                if is_client_feature and self.fit_select_features:
-                    self.zero_shap_client_features.append(original_name)
                 continue
 
             # Use only important features
@@ -4525,6 +4413,9 @@ if response.status_code == 200:
             features_info_without_links.append(feature_info.to_row_without_links(self.bundle))
             internal_features_info.append(feature_info.to_internal_row(self.bundle))
 
+        if update_selected_features:
+            self._search_task.update_selected_features(trace_id, self.feature_names_)
+
         if len(features_info) > 0:
             self.features_info = pd.DataFrame(features_info)
             if self.features_info[self.bundle.get("features_info_psi")].isna().all():
@@ -4652,32 +4543,10 @@ if response.status_code == 200:
             )
         )
 
-    def __filtered_importance_names(
-        self, importance_threshold: Optional[float], max_features: Optional[int], trace_id: str, df: pd.DataFrame
-    ) -> List[str]:
-        # get features importance from server
-        filtered_importances = self.__get_features_importance_from_server(trace_id, df)
-
-        if len(filtered_importances) == 0:
-            return []
-
-        if importance_threshold is not None:
-            filtered_importances = [
-                (name, importance)
-                for name, importance in filtered_importances.items()
-                if importance > importance_threshold
-            ]
-        if max_features is not None:
-            filtered_importances = list(filtered_importances)[:max_features]
-        if len(filtered_importances) == 0:
-            return []
-        filtered_importance_names, _ = zip(*filtered_importances)
-        return list(filtered_importance_names)
-
     def __prepare_search_keys(
         self,
         x: pd.DataFrame,
-        search_keys: Dict[str, SearchKey],
+        search_keys: dict[str, SearchKey],
         is_demo_dataset: bool,
         is_transform=False,
         silent_mode=False,
@@ -4788,20 +4657,16 @@ if response.status_code == 200:
 
     def __show_metrics(
         self,
-        scoring: Union[Callable, str, None],
-        estimator: Optional[Any],
-        importance_threshold: Optional[float],
-        max_features: Optional[int],
-        remove_outliers_calc_metrics: Optional[bool],
+        scoring: Callable | str | None,
+        estimator: Any | None,
+        remove_outliers_calc_metrics: bool | None,
         trace_id: str,
-        progress_bar: Optional[ProgressBar] = None,
-        progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
+        progress_bar: ProgressBar | None = None,
+        progress_callback: Callable[[SearchProgress], Any] | None = None,
     ):
         self.metrics = self.calculate_metrics(
             scoring=scoring,
             estimator=estimator,
-            importance_threshold=importance_threshold,
-            max_features=max_features,
             remove_outliers_calc_metrics=remove_outliers_calc_metrics,
             trace_id=trace_id,
             internal_call=True,
@@ -4812,22 +4677,15 @@ if response.status_code == 200:
             msg = self.bundle.get("quality_metrics_header")
             display_html_dataframe(self.metrics, self.metrics, msg)
 
-    def __show_selected_features(self, search_keys: Dict[str, SearchKey]):
-        search_key_names = [col for col, tpe in search_keys.items() if tpe != SearchKey.CUSTOM_KEY]
-        if self.fit_columns_renaming:
-            search_key_names = sorted(set([self.fit_columns_renaming.get(col, col) for col in search_key_names]))
-        msg = self.bundle.get("features_info_header").format(len(self.feature_names_), search_key_names)
-
+    def __show_selected_features(self):
         try:
             _ = get_ipython()  # type: ignore
 
-            print(Format.GREEN + Format.BOLD + msg + Format.END)
-            self.logger.info(msg)
             if len(self.feature_names_) > 0:
                 self.features_info_display_handle = display_html_dataframe(
                     self.features_info,
                     self._features_info_without_links,
-                    self.bundle.get("relevant_features_header"),
+                    self.bundle.get("relevant_features_header").format(len(self.feature_names_)),
                     display_id=f"features_info_{uuid.uuid4()}",
                 )
 
@@ -4838,14 +4696,23 @@ if response.status_code == 200:
                         self.bundle.get("relevant_data_sources_header"),
                         display_id=f"data_sources_{uuid.uuid4()}",
                     )
+
+                autofe_description = self.get_autofe_features_description()
+                if autofe_description is not None and len(autofe_description) > 0:
+                    self.logger.info(f"AutoFE descriptions: {autofe_description}")
+                    self.autofe_features_display_handle = display_html_dataframe(
+                        df=autofe_description,
+                        internal_df=autofe_description,
+                        header=self.bundle.get("autofe_descriptions_header"),
+                        display_id=f"autofe_descriptions_{uuid.uuid4()}",
+                    )
             else:
                 msg = self.bundle.get("features_info_zero_important_features")
                 self.__log_warning(msg, show_support_link=True)
         except (ImportError, NameError):
-            print(msg)
             print(self._internal_features_info)
 
-    def __show_report_button(self, display_id: Optional[str] = None, display_handle=None):
+    def __show_report_button(self, display_id: str | None = None, display_handle=None):
         try:
             return prepare_and_show_report(
                 relevant_features_df=self._features_info_without_links,
@@ -4861,40 +4728,14 @@ if response.status_code == 200:
         except Exception:
             pass
 
-    def __validate_importance_threshold(self, importance_threshold: Optional[float]) -> float:
-        try:
-            return float(importance_threshold) if importance_threshold is not None else 0.0
-        except ValueError:
-            self.logger.exception(f"Invalid importance_threshold provided: {importance_threshold}")
-            raise ValidationError(self.bundle.get("invalid_importance_threshold"))
-
-    def __validate_max_features(self, max_features: Optional[int]) -> int:
-        try:
-            return int(max_features) if max_features is not None else 400
-        except ValueError:
-            self.logger.exception(f"Invalid max_features provided: {max_features}")
-            raise ValidationError(self.bundle.get("invalid_max_features"))
-
-    def __filtered_enriched_features(
-        self,
-        importance_threshold: Optional[float],
-        max_features: Optional[int],
-        trace_id: str,
-        df: pd.DataFrame,
-    ) -> List[str]:
-        importance_threshold = self.__validate_importance_threshold(importance_threshold)
-        max_features = self.__validate_max_features(max_features)
-
-        return self.__filtered_importance_names(importance_threshold, max_features, trace_id, df)
-
     def __detect_missing_search_keys(
         self,
         df: pd.DataFrame,
-        search_keys: Dict[str, SearchKey],
+        search_keys: dict[str, SearchKey],
         is_demo_dataset: bool,
         silent_mode=False,
         is_transform=False,
-    ) -> Dict[str, SearchKey]:
+    ) -> dict[str, SearchKey]:
         sample = df.head(100)
 
         def check_need_detect(search_key: SearchKey):
@@ -5011,7 +4852,7 @@ if response.status_code == 200:
         except Exception:
             self.logger.exception("Failed to dump python libs")
 
-    def __display_support_link(self, link_text: Optional[str] = None):
+    def __display_support_link(self, link_text: str | None = None):
         support_link = self.bundle.get("support_link")
         link_text = link_text or self.bundle.get("support_text")
         try:
@@ -5038,9 +4879,9 @@ if response.status_code == 200:
     def dump_input(
         self,
         trace_id: str,
-        X: Union[pd.DataFrame, pd.Series],
-        y: Union[pd.DataFrame, pd.Series, None] = None,
-        eval_set: Union[Tuple, None] = None,
+        X: pd.DataFrame | pd.Series,
+        y: pd.DataFrame | pd.Series | None = None,
+        eval_set: tuple | None = None,
     ):
         def dump_task(X_, y_, eval_set_):
             with MDC(trace_id=trace_id):
@@ -5131,28 +4972,10 @@ def is_frames_equal(first, second, bundle: ResourceBundle) -> bool:
         raise ValidationError(bundle.get("x_and_eval_x_diff_types").format(type(first), type(second)))
 
 
-def drop_duplicates(df: Union[pd.DataFrame, np.ndarray]) -> pd.DataFrame:
+def drop_duplicates(df: pd.DataFrame | np.ndarray | Any) -> pd.DataFrame:
     if isinstance(df, pd.DataFrame):
         return df.drop_duplicates()
     elif isinstance(df, np.ndarray):
         return pd.DataFrame(df).drop_duplicates()
     else:
         return df
-
-
-def hash_input(X: pd.DataFrame, y: Optional[pd.Series] = None, eval_set: Optional[List[Tuple]] = None) -> str:
-    hashed_objects = []
-    try:
-        hashed_objects.append(pd.util.hash_pandas_object(X, index=False).values)
-        if y is not None:
-            hashed_objects.append(pd.util.hash_pandas_object(y, index=False).values)
-        if eval_set is not None:
-            if isinstance(eval_set, tuple):
-                eval_set = [eval_set]
-            for eval_X, eval_y in eval_set:
-                hashed_objects.append(pd.util.hash_pandas_object(eval_X, index=False).values)
-                hashed_objects.append(pd.util.hash_pandas_object(eval_y, index=False).values)
-        common_hash = hashlib.sha256(np.concatenate(hashed_objects)).hexdigest()
-        return common_hash
-    except Exception:
-        return ""