upgini 1.2.113a1__py3-none-any.whl → 1.2.113a2__py3-none-any.whl

upgini/features_enricher.py

@@ -112,6 +112,7 @@ except Exception:
         CustomFallbackProgressBar as ProgressBar,
     )
 
+from upgini.utils.psi import calculate_features_psi
 from upgini.utils.sample_utils import SampleColumns, SampleConfig, _num_samples, sample
 from upgini.utils.sort import sort_columns
 from upgini.utils.target_utils import calculate_psi, define_task
@@ -297,7 +298,9 @@ class FeaturesEnricher(TransformerMixin):
         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.search_id = search_id
         self.disable_force_downsampling = disable_force_downsampling
         self.print_trace_id = print_trace_id
@@ -398,13 +401,26 @@ class FeaturesEnricher(TransformerMixin):
     @staticmethod
     def _check_eval_set(eval_set, X, bundle: ResourceBundle):
         checked_eval_set = []
-        if eval_set is not None and isinstance(eval_set, tuple):
+        if eval_set is None:
+            return checked_eval_set
+        if isinstance(eval_set, tuple):
             eval_set = [eval_set]
-        if eval_set is not None and not isinstance(eval_set, list):
+        if not isinstance(eval_set, list):
             raise ValidationError(bundle.get("unsupported_type_eval_set").format(type(eval_set)))
         for eval_pair in eval_set or []:
+            # Handle OOT
+            if isinstance(eval_pair, pd.DataFrame):
+                empty_target = pd.Series([np.nan] * len(eval_pair), index=eval_pair.index)
+                eval_pair = (eval_pair, empty_target)
+            elif isinstance(eval_pair, tuple) and len(eval_pair) == 1:
+                empty_target = pd.Series([np.nan] * len(eval_pair[0]), index=eval_pair[0].index)
+                eval_pair = (eval_pair[0], empty_target)
+
             if not isinstance(eval_pair, tuple) or len(eval_pair) != 2:
                 raise ValidationError(bundle.get("eval_set_invalid_tuple_size").format(len(eval_pair)))
+            if eval_pair[1] is None:
+                empty_target = pd.Series([np.nan] * len(eval_pair[0]), index=eval_pair[0].index)
+                eval_pair = (eval_pair[0], empty_target)
             if not is_frames_equal(X, eval_pair[0], bundle):
                 checked_eval_set.append(eval_pair)
         return checked_eval_set
@@ -415,7 +431,6 @@ class FeaturesEnricher(TransformerMixin):
         y: Union[pd.Series, np.ndarray, List],
         eval_set: Optional[Union[List[tuple], tuple]] = None,
         *args,
-        oot: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
         exclude_features_sources: Optional[List[str]] = None,
         calculate_metrics: Optional[bool] = None,
         estimator: Optional[Any] = None,
@@ -427,6 +442,7 @@ class FeaturesEnricher(TransformerMixin):
         search_id_callback: Optional[Callable[[str], Any]] = None,
         select_features: bool = True,
         auto_fe_parameters: Optional[AutoFEParameters] = None,
+        stability_threshold: float = 0.15,
         **kwargs,
     ):
         """Fit to data.
@@ -444,9 +460,6 @@ class FeaturesEnricher(TransformerMixin):
         eval_set: List[tuple], optional (default=None)
             List of pairs (X, y) for validation.
 
-        oot: pandas.DataFrame of shape (n_samples, n_features)
-            Out of time data.
-
         importance_threshold: float, optional (default=None)
             Minimum SHAP value to select a feature. Default value is 0.0.
 
@@ -512,14 +525,14 @@ class FeaturesEnricher(TransformerMixin):
                     X,
                     y,
                     self.eval_set,
-                    oot=oot,
-                    progress_bar=progress_bar,
+                    progress_bar,
                     start_time=start_time,
                     exclude_features_sources=exclude_features_sources,
                     calculate_metrics=calculate_metrics,
                     estimator=estimator,
                     scoring=scoring,
                     importance_threshold=importance_threshold,
+                    stability_threshold=stability_threshold,
                     max_features=max_features,
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
                     auto_fe_parameters=auto_fe_parameters,
@@ -568,7 +581,6 @@ class FeaturesEnricher(TransformerMixin):
         y: Union[pd.DataFrame, pd.Series, np.ndarray, List],
         eval_set: Optional[Union[List[tuple], tuple]] = None,
         *args,
-        oot: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
         exclude_features_sources: Optional[List[str]] = None,
         keep_input: bool = True,
         importance_threshold: Optional[float] = None,
@@ -580,6 +592,7 @@ class FeaturesEnricher(TransformerMixin):
         progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
         select_features: bool = True,
         auto_fe_parameters: Optional[AutoFEParameters] = None,
+        stability_threshold: float = 0.15,
         **kwargs,
     ) -> pd.DataFrame:
         """Fit to data, then transform it.
@@ -624,6 +637,10 @@ 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 for selected features PSI calculation. If PSI is less than this threshold,
+            then feature will be dropped.
+
         Returns
         -------
         X_new: pandas.DataFrame of shape (n_samples, n_features_new)
@@ -673,14 +690,14 @@ class FeaturesEnricher(TransformerMixin):
                     X,
                     y,
                     self.eval_set,
-                    oot=oot,
-                    progress_bar=progress_bar,
+                    progress_bar,
                     start_time=start_time,
                     exclude_features_sources=exclude_features_sources,
                     calculate_metrics=calculate_metrics,
                     scoring=scoring,
                     estimator=estimator,
                     importance_threshold=importance_threshold,
+                    stability_threshold=stability_threshold,
                     max_features=max_features,
                     remove_outliers_calc_metrics=remove_outliers_calc_metrics,
                     auto_fe_parameters=auto_fe_parameters,
@@ -947,8 +964,8 @@ class FeaturesEnricher(TransformerMixin):
             ):
                 raise ValidationError(self.bundle.get("metrics_unfitted_enricher"))
 
-            validated_X, validated_y, validated_eval_set, _ = self._validate_train_eval(
-                effective_X, effective_y, effective_eval_set
+            validated_X, validated_y, validated_eval_set = self._validate_train_eval(
+                effective_X, effective_y, effective_eval_set, silent=internal_call
             )
 
             if self.X is None:
@@ -1040,6 +1057,7 @@ class FeaturesEnricher(TransformerMixin):
                     groups,
                     _cv,
                     columns_renaming,
+                    eval_set_dates,
                 ) = prepared_data
 
                 # rename cat_features
@@ -1081,9 +1099,9 @@ class FeaturesEnricher(TransformerMixin):
                     self._check_train_and_eval_target_distribution(y_sorted, fitting_eval_set_dict)
 
                     has_date = self._get_date_column(search_keys) is not None
-                    has_time = has_date and isinstance(_cv, TimeSeriesSplit) or isinstance(_cv, BlockedTimeSeriesSplit)
                     model_task_type = self.model_task_type or define_task(y_sorted, has_date, self.logger, silent=True)
                     cat_features = list(set(client_cat_features + cat_features_from_backend))
+                    has_time = has_date and isinstance(_cv, TimeSeriesSplit) or isinstance(_cv, BlockedTimeSeriesSplit)
                     baseline_cat_features = [f for f in cat_features if f in fitting_X.columns]
                     enriched_cat_features = [f for f in cat_features if f in fitting_enriched_X.columns]
                     if len(enriched_cat_features) < len(cat_features):
@@ -1203,8 +1221,6 @@ class FeaturesEnricher(TransformerMixin):
                     # max_initial_eval_set_hit_rate = self._search_task.get_max_initial_eval_set_hit_rate_v2()
                     if len(fitting_eval_set_dict) > 0:
                         for idx in fitting_eval_set_dict.keys():
-                            # eval_hit_rate = max_initial_eval_set_hit_rate[idx + 1]
-
                             (
                                 eval_X_sorted,
                                 eval_y_sorted,
@@ -1212,6 +1228,10 @@ class FeaturesEnricher(TransformerMixin):
                                 enriched_eval_y_sorted,
                             ) = fitting_eval_set_dict[idx]
 
+                            if eval_y_sorted.isna().all():
+                                # Skip OOT eval set
+                                continue
+
                             if baseline_estimator is not None:
                                 self.logger.info(
                                     f"Calculate baseline {metric} on eval set {idx + 1} "
@@ -1254,17 +1274,14 @@ class FeaturesEnricher(TransformerMixin):
                                     "quality_metrics_eval_segment"
                                 ).format(idx + 1),
                                 self.bundle.get("quality_metrics_rows_header"): _num_samples(
-                                    # effective_eval_set[idx][0]
                                     # Use actually used for metrics dataset
                                     eval_X_sorted
                                 ),
-                                # self.bundle.get("quality_metrics_match_rate_header"): eval_hit_rate,
                             }
                             if model_task_type in [ModelTaskType.BINARY, ModelTaskType.REGRESSION] and is_numeric_dtype(
                                 eval_y_sorted
                             ):
                                 eval_metrics[self.bundle.get("quality_metrics_mean_target_header")] = round(
-                                    # np.mean(validated_eval_set[idx][1]), 4
                                     # Use actually used for metrics dataset
                                     np.mean(eval_y_sorted),
                                     4,
@@ -1337,6 +1354,199 @@ class FeaturesEnricher(TransformerMixin):
             finally:
                 self.logger.info(f"Calculating metrics elapsed time: {time.time() - start_time}")
 
+    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]],
+        stability_threshold: float,
+        cv: Union[BaseCrossValidator, CVType, str, None] = None,
+        estimator=None,
+        exclude_features_sources: Optional[List[str]] = None,
+        importance_threshold: Optional[float] = None,
+        max_features: Optional[int] = None,
+        progress_bar: bool = True,
+        progress_callback: Optional[Callable] = None,
+    ):
+        search_keys = self.search_keys.copy()
+        validated_X, _, validated_eval_set = self._validate_train_eval(X, y, eval_set, silent=True)
+        if isinstance(X, np.ndarray):
+            search_keys = {str(k): v for k, v in search_keys.items()}
+
+        has_date = self._get_date_column(search_keys) is not None
+        if not has_date or not validated_eval_set:
+            self.logger.info("No date column or eval set for OOT psi calculation")
+            return
+
+        cat_features_from_backend = self.__get_categorical_features()
+        client_cat_features, search_keys_for_metrics = self._get_and_validate_client_cat_features(
+            estimator, validated_X, search_keys
+        )
+        if self.id_columns and self.id_columns_encoder is not None:
+            if cat_features_from_backend:
+                cat_features_from_backend = [
+                    c
+                    for c in cat_features_from_backend
+                    if self.fit_columns_renaming.get(c, c) not in self.id_columns_encoder.feature_names_in_
+                ]
+            if client_cat_features:
+                client_cat_features = [
+                    c
+                    for c in client_cat_features
+                    if self.fit_columns_renaming.get(c, c) not in self.id_columns_encoder.feature_names_in_
+                ]
+
+        prepared_data = self._prepare_data_for_metrics(
+            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,
+            progress_bar=progress_bar,
+            progress_callback=progress_callback,
+            client_cat_features=client_cat_features,
+        )
+        if prepared_data is None:
+            return None
+
+        (
+            validated_X,
+            fitting_X,
+            y_sorted,
+            fitting_enriched_X,
+            _,
+            fitting_eval_set_dict,
+            _,
+            _,
+            _,
+            columns_renaming,
+            eval_set_dates,
+        ) = prepared_data
+
+        # rename cat_features
+        if client_cat_features:
+            for new_c, old_c in columns_renaming.items():
+                if old_c in client_cat_features:
+                    client_cat_features.remove(old_c)
+                    client_cat_features.append(new_c)
+            for cat_feature in client_cat_features:
+                if cat_feature not in fitting_X.columns:
+                    self.logger.error(
+                        f"Client cat_feature `{cat_feature}` not found in" f" x columns: {fitting_X.columns.to_list()}"
+                    )
+        else:
+            client_cat_features = []
+
+        model_task_type = self.model_task_type or define_task(y_sorted, has_date, self.logger, silent=True)
+        cat_features = list(set(client_cat_features + cat_features_from_backend))
+
+        # Drop unstable features
+        unstable_features = self._check_stability(
+            validated_X,
+            validated_eval_set,
+            fitting_eval_set_dict,
+            eval_set_dates,
+            search_keys,
+            stability_threshold,
+            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"
+            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],
+        stability_threshold: float,
+        cat_features: List[str],
+        model_task_type: ModelTaskType,
+    ) -> List[str]:
+        # Find latest eval set or earliest if all eval sets are before train set
+        date_column = self._get_date_column(search_keys)
+
+        if (
+            date_column is None
+            or not eval_set
+            or not eval_set_dates
+            or (self.cv is not None and self.cv.is_time_series())
+        ):
+            return []
+
+        # Get minimum date from main dataset X
+        main_min_date = X[date_column].min()
+
+        # Find minimum date for each eval_set and compare with main dataset
+        eval_dates = []
+        for i, (eval_x, _) in enumerate(eval_set):
+            if date_column in eval_x.columns:
+                eval_min_date = eval_x[date_column].min()
+                eval_max_date = eval_x[date_column].max()
+                eval_dates.append((i, eval_min_date, eval_max_date))
+
+        if not eval_dates:
+            return []
+
+        # Check if any eval_set has minimum date >= main dataset minimum date
+        later_eval_sets = [(i, min_date, max_date) for i, min_date, max_date in eval_dates if min_date >= main_min_date]
+
+        if later_eval_sets:
+            # If there are eval_sets with date >= main date, choose the one with highest maximum date
+            selected_eval_set_idx = max(later_eval_sets, key=lambda x: x[2])[0]
+        else:
+            # If all eval_sets have dates < main date, choose the one with lowest minimux date
+            selected_eval_set_idx = max(eval_dates, key=lambda x: x[1])[0]
+
+        checking_eval_set = enriched_eval_set[selected_eval_set_idx]
+
+        checking_eval_set_df = (
+            checking_eval_set[2]
+            if checking_eval_set[1] is None or checking_eval_set[1].isna().all()
+            else pd.concat([checking_eval_set[2], checking_eval_set[1].to_frame(TARGET)], axis=1)
+        )
+        checking_eval_set_df = checking_eval_set_df.copy()
+
+        checking_eval_set_df[date_column] = eval_set_dates[selected_eval_set_idx]
+
+        psi_values = calculate_features_psi(
+            checking_eval_set_df, cat_features, date_column, self.logger, model_task_type
+        )
+
+        self.psi_values = {
+            feature: psi_value for feature, psi_value in psi_values.items() if psi_value <= stability_threshold
+        }
+
+        return [feature for feature, psi in psi_values.items() if psi > stability_threshold]
+
     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}")
@@ -1392,6 +1602,56 @@ 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():
@@ -1542,7 +1802,7 @@ class FeaturesEnricher(TransformerMixin):
         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, self.bundle)
-        validated_X, validated_y, validated_eval_set, _ = self._validate_train_eval(X, y, checked_eval_set)
+        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,
@@ -1556,7 +1816,7 @@ class FeaturesEnricher(TransformerMixin):
             progress_bar,
             progress_callback,
         )
-        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming = dataclasses.astuple(
+        (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys, columns_renaming) = dataclasses.astuple(
             sampled_data
         )
 
@@ -1579,8 +1839,7 @@ class FeaturesEnricher(TransformerMixin):
                 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 (
+            and c not in (
                 excluding_search_keys
                 + list(self.fit_dropped_features)
                 + [DateTimeSearchKeyConverter.DATETIME_COL, SYSTEM_RECORD_ID, ENTITY_SYSTEM_RECORD_ID]
@@ -1679,12 +1938,16 @@ class FeaturesEnricher(TransformerMixin):
         fitting_enriched_X = fitting_enriched_X[fitting_enriched_x_columns]
         fitting_enriched_X, _ = self._encode_id_columns(fitting_enriched_X, self.fit_columns_renaming)
         self.logger.info(f"Final sorted list of fitting enriched X columns: {fitting_enriched_x_columns}")
+        date_column = self._get_date_column(search_keys)
+        eval_set_dates = {}
         for idx, eval_tuple in eval_set_sampled_dict.items():
             eval_X_sampled, enriched_eval_X, eval_y_sampled = eval_tuple
             eval_X_sorted, eval_y_sorted = self._sort_by_system_record_id(eval_X_sampled, eval_y_sampled, self.cv)
             enriched_eval_X_sorted, enriched_eval_y_sorted = self._sort_by_system_record_id(
                 enriched_eval_X, eval_y_sampled, self.cv
             )
+            if date_column is not None:
+                eval_set_dates[idx] = eval_X_sorted[date_column]
             fitting_eval_X = eval_X_sorted[fitting_x_columns].copy()
             fitting_enriched_eval_X = enriched_eval_X_sorted[fitting_enriched_x_columns].copy()
 
@@ -1729,6 +1992,7 @@ class FeaturesEnricher(TransformerMixin):
             groups,
             cv,
             columns_renaming,
+            eval_set_dates,
         )
 
     @dataclass
@@ -1891,14 +2155,16 @@ class FeaturesEnricher(TransformerMixin):
         remove_outliers_calc_metrics: Optional[bool],
     ) -> _EnrichedDataForMetrics:
         eval_set_sampled_dict = {}
-        search_keys = self.fit_search_keys
+        search_keys = self.fit_search_keys.copy()
 
         rows_to_drop = None
         has_date = self._get_date_column(search_keys) is not None
         self.model_task_type = self.model_task_type or define_task(
             self.df_with_original_index[TARGET], has_date, self.logger, silent=True
         )
-        if self.model_task_type == ModelTaskType.REGRESSION:
+        if remove_outliers_calc_metrics is None:
+            remove_outliers_calc_metrics = True
+        if self.model_task_type == ModelTaskType.REGRESSION and remove_outliers_calc_metrics:
             target_outliers_df = self._search_task.get_target_outliers(trace_id)
             if target_outliers_df is not None and len(target_outliers_df) > 0:
                 outliers = pd.merge(
@@ -1908,11 +2174,8 @@ class FeaturesEnricher(TransformerMixin):
                     how="inner",
                 )
                 top_outliers = outliers.sort_values(by=TARGET, ascending=False)[TARGET].head(3)
-                if remove_outliers_calc_metrics is None or remove_outliers_calc_metrics is True:
-                    rows_to_drop = outliers
-                    not_msg = ""
-                else:
-                    not_msg = "not "
+                rows_to_drop = outliers
+                not_msg = ""
                 msg = self.bundle.get("target_outliers_warning").format(len(target_outliers_df), top_outliers, not_msg)
                 print(msg)
                 self.logger.warning(msg)
@@ -1938,8 +2201,11 @@ class FeaturesEnricher(TransformerMixin):
         )
 
         # Handle eval sets extraction based on EVAL_SET_INDEX
-        if EVAL_SET_INDEX in enriched_Xy.columns and eval_set is not None:
-            for eval_set_index in range(1, len(eval_set) + 1):
+        if EVAL_SET_INDEX in enriched_Xy.columns:
+            eval_set_indices = list(enriched_Xy[EVAL_SET_INDEX].unique())
+            if 0 in eval_set_indices:
+                eval_set_indices.remove(0)
+            for eval_set_index in eval_set_indices:
                 enriched_eval_sets[eval_set_index] = enriched_Xy.loc[
                     enriched_Xy[EVAL_SET_INDEX] == eval_set_index
                 ].copy()
@@ -1967,12 +2233,13 @@ class FeaturesEnricher(TransformerMixin):
                 enriched_eval_X = enriched_eval_sets[idx + 1][enriched_X_columns].copy()
                 eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
 
-        reversed_renaming = {v: k for k, v in self.fit_columns_renaming.items()}
-        X_sampled.rename(columns=reversed_renaming, inplace=True)
-        enriched_X.rename(columns=reversed_renaming, inplace=True)
+        # reversed_renaming = {v: k for k, v in self.fit_columns_renaming.items()}
+        X_sampled.rename(columns=self.fit_columns_renaming, inplace=True)
+        enriched_X.rename(columns=self.fit_columns_renaming, inplace=True)
         for _, (eval_X_sampled, enriched_eval_X, _) in eval_set_sampled_dict.items():
-            eval_X_sampled.rename(columns=reversed_renaming, inplace=True)
-            enriched_eval_X.rename(columns=reversed_renaming, inplace=True)
+            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()}
 
         datasets_hash = hash_input(self.X, self.y, self.eval_set)
         return self.__cache_and_return_results(
@@ -2051,11 +2318,7 @@ class FeaturesEnricher(TransformerMixin):
         )
 
     def __combine_train_and_eval_sets(
-        self,
-        X: pd.DataFrame,
-        y: Optional[pd.Series] = None,
-        eval_set: Optional[List[tuple]] = None,
-        oot: Optional[pd.DataFrame] = None,
+        self, X: pd.DataFrame, y: Optional[pd.Series] = None, eval_set: Optional[List[tuple]] = None
     ) -> pd.DataFrame:
         df = X.copy()
         if y is not None:
@@ -2071,11 +2334,6 @@ class FeaturesEnricher(TransformerMixin):
             eval_df_with_index[TARGET] = eval_y
             eval_df_with_index[EVAL_SET_INDEX] = idx + 1
             df = pd.concat([df, eval_df_with_index])
-        
-        if oot is not None:
-            oot_df_with_index = oot.copy()
-            oot_df_with_index[EVAL_SET_INDEX] = -1
-            df = pd.concat([df, oot_df_with_index])
 
         return df
 
@@ -2125,15 +2383,15 @@ class FeaturesEnricher(TransformerMixin):
 
     def __extract_eval_data(
         self, enriched_df: pd.DataFrame, x_columns: List[str], enriched_X_columns: List[str], eval_set_len: int
-    ) -> Dict[int, Tuple]:
+    ) -> Tuple[Dict[int, Tuple], Dict[int, pd.Series]]:
         eval_set_sampled_dict = {}
 
-        for idx in range(1, eval_set_len + 1):
-            enriched_eval_xy = enriched_df.query(f"{EVAL_SET_INDEX} == {idx}")
+        for idx in range(eval_set_len):
+            enriched_eval_xy = enriched_df.query(f"{EVAL_SET_INDEX} == {idx + 1}")
             eval_x_sampled = enriched_eval_xy[x_columns].copy()
             eval_y_sampled = enriched_eval_xy[TARGET].copy()
             enriched_eval_x = enriched_eval_xy[enriched_X_columns].copy()
-            eval_set_sampled_dict[idx - 1] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
+            eval_set_sampled_dict[idx] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
 
         return eval_set_sampled_dict
 
@@ -2171,12 +2429,12 @@ class FeaturesEnricher(TransformerMixin):
         columns_renaming: Dict[str, 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()
-        }
+        # 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,
@@ -2325,10 +2583,10 @@ if response.status_code == 200:
         with MDC(trace_id=trace_id, search_id=search_id):
             self.logger.info("Start transform")
 
-            validated_X, validated_y, _, _ = self._validate_train_eval(
-                X, y, is_transform=True
+            validated_X, validated_y, validated_eval_set = self._validate_train_eval(
+                X, y, eval_set=None, is_transform=True, silent=True
             )
-            df = self.__combine_train_and_eval_sets(validated_X, validated_y)
+            df = self.__combine_train_and_eval_sets(validated_X, validated_y, validated_eval_set)
 
             validated_Xy = df.copy()
 
@@ -2428,7 +2686,7 @@ if response.status_code == 200:
             else:
                 self.logger.info("Input dataset hasn't date column")
                 if self.__should_add_date_column():
-                    df = self._add_current_date_as_key(df, search_keys, self.logger, self.bundle)
+                    df = self._add_current_date_as_key(df, search_keys, self.bundle, silent=True)
 
             email_columns = SearchKey.find_all_keys(search_keys, SearchKey.EMAIL)
             if email_columns and self.generate_search_key_features:
@@ -2677,7 +2935,8 @@ if response.status_code == 200:
             selecting_columns = [
                 c
                 for c in itertools.chain(validated_Xy.columns.tolist(), selected_generated_features)
-                if c not in self.zero_shap_client_features or c in (self.id_columns or [])
+                if (c not in self.zero_shap_client_features and c not in self.unstable_client_features)
+                or c in (self.id_columns or [])
             ]
             selecting_columns.extend(c for c in result.columns if c in filtered_columns and c not in selecting_columns)
             if add_fit_system_record_id:
@@ -2803,15 +3062,15 @@ if response.status_code == 200:
         X: Union[pd.DataFrame, pd.Series, np.ndarray],
         y: Union[pd.DataFrame, pd.Series, np.ndarray, List, None],
         eval_set: Optional[List[tuple]],
-        *,
-        oot: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
         progress_bar: Optional[ProgressBar],
         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],
+        stability_threshold: float,
         max_features: Optional[int],
         remove_outliers_calc_metrics: Optional[bool],
         auto_fe_parameters: AutoFEParameters,
@@ -2826,8 +3085,9 @@ if response.status_code == 200:
         self.fit_columns_renaming = None
         self.fit_dropped_features = set()
         self.fit_generated_features = []
+        self.psi_values = None
 
-        validated_X, validated_y, validated_eval_set, validated_oot = self._validate_train_eval(X, y, eval_set, oot)
+        validated_X, validated_y, validated_eval_set = self._validate_train_eval(X, y, eval_set)
 
         is_demo_dataset = hash_input(validated_X, validated_y, validated_eval_set) in DEMO_DATASET_HASHES
         if is_demo_dataset:
@@ -2868,7 +3128,6 @@ if response.status_code == 200:
             validated_X,
             validated_y,
             validated_eval_set,
-            validated_oot,
             exclude_features_sources=exclude_features_sources,
             calculate_metrics=calculate_metrics,
             scoring=scoring,
@@ -2876,7 +3135,7 @@ if response.status_code == 200:
             remove_outliers_calc_metrics=remove_outliers_calc_metrics,
         )
 
-        df = self.__combine_train_and_eval_sets(validated_X, validated_y, validated_eval_set, validated_oot)
+        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()
@@ -2923,7 +3182,7 @@ if response.status_code == 200:
             self.logger.info("Input dataset hasn't date column")
             # TODO remove when this logic will be implemented on the back
             if self.__should_add_date_column():
-                df = self._add_current_date_as_key(df, self.fit_search_keys, self.logger, self.bundle)
+                df = self._add_current_date_as_key(df, self.fit_search_keys, self.bundle)
 
         email_columns = SearchKey.find_all_keys(self.fit_search_keys, SearchKey.EMAIL)
         if email_columns and self.generate_search_key_features:
@@ -2938,10 +3197,13 @@ if response.status_code == 200:
         except Exception:
             self.logger.exception("Failed to check dates distribution validity")
 
+        self.__adjust_cv(df)
+
         if (
             is_numeric_dtype(df[self.TARGET_NAME])
             and self.model_task_type in [ModelTaskType.BINARY, ModelTaskType.MULTICLASS]
             and has_date
+            and not self.cv.is_time_series()
         ):
             self._validate_PSI(df.sort_values(by=maybe_date_column))
 
@@ -2973,8 +3235,9 @@ if response.status_code == 200:
 
         self.fit_generated_features = [f for f in self.fit_generated_features if f not in self.fit_dropped_features]
 
+        # Group columns should have normalized names
+        self.cv = None
         self.__adjust_cv(df)
-
         if self.id_columns is not None and self.cv is not None and self.cv.is_time_series():
             id_columns = self.__get_renamed_id_columns()
             if id_columns:
@@ -3179,6 +3442,21 @@ if response.status_code == 200:
                     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,
+                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,
+            )
+
             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")
@@ -3303,14 +3581,13 @@ if response.status_code == 200:
         X: pd.DataFrame,
         y: Optional[pd.Series] = None,
         eval_set: Optional[List[Tuple[pd.DataFrame, pd.Series]]] = None,
-        oot: Union[pd.DataFrame, pd.Series, np.ndarray, None] = None,
         is_transform: bool = False,
+        silent: bool = False,
     ) -> Tuple[pd.DataFrame, pd.Series, Optional[List[Tuple[pd.DataFrame, pd.Series]]]]:
         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)
-        validated_oot = self._validate_oot(validated_X, oot)
-        return validated_X, validated_y, validated_eval_set, validated_oot
+        validated_eval_set = self._validate_eval_set(validated_X, eval_set, silent=silent)
+        return validated_X, validated_y, validated_eval_set
 
     def _encode_id_columns(
         self,
@@ -3441,53 +3718,30 @@ 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]]]):
+    def _validate_eval_set(
+        self, X: pd.DataFrame, eval_set: Optional[List[Tuple[pd.DataFrame, pd.Series]]], silent: bool = False
+    ):
         if eval_set is None:
             return None
-        return [self._validate_eval_set_pair(X, eval_pair) for eval_pair in eval_set]
-
-    def _validate_oot(self, X: pd.DataFrame, oot: Optional[pd.DataFrame]):
-        if oot is None:
-            return None
-
-        if _num_samples(oot) == 0:
-            raise ValidationError(self.bundle.get("oot_is_empty"))
-        if isinstance(oot, pd.DataFrame):
-            if isinstance(oot.columns, pd.MultiIndex) or isinstance(oot.index, pd.MultiIndex):
-                raise ValidationError(self.bundle.get("oot_multiindex_unsupported"))
-            validated_oot = oot.copy()
-        elif isinstance(oot, pd.Series):
-            validated_oot = oot.to_frame()
-        elif isinstance(oot, (list, np.ndarray)):
-            validated_oot = pd.DataFrame(oot)
-            renaming = {c: str(c) for c in validated_oot.columns}
-            validated_oot = validated_oot.rename(columns=renaming)
-        else:
-            raise ValidationError(self.bundle.get("unsupported_type_oot").format(type(oot)))
-
-        if not validated_oot.index.is_unique:
-            raise ValidationError(self.bundle.get("non_unique_index_oot"))
-
-        if self.exclude_columns is not None:
-            validated_oot = validated_oot.drop(columns=self.exclude_columns, errors="ignore")
-
-        if self.baseline_score_column:
-            validated_oot[self.baseline_score_column] = validated_oot[self.baseline_score_column].astype(
-                "float64", errors="ignore"
-            )
-
-        if validated_oot.columns.to_list() != X.columns.to_list():
-            if set(validated_oot.columns.to_list()) == set(X.columns.to_list()):
-                validated_oot = validated_oot[X.columns.to_list()]
-            else:
-                raise ValidationError(self.bundle.get("oot_and_x_diff_shape"))
-
-        # Check for duplicates between train and eval sets by comparing all values
-        train_eval_intersection = pd.merge(X, validated_oot, how="inner")
-        if len(train_eval_intersection) > 0:
-            raise ValidationError(self.bundle.get("oot_has_train_samples"))
+        validated_eval_set = []
+        has_date = self._get_date_column(self.search_keys) is not None
+        for idx, eval_pair in enumerate(eval_set):
+            validated_pair = self._validate_eval_set_pair(X, eval_pair)
+            if validated_pair[1].isna().all():
+                if not has_date:
+                    msg = self.bundle.get("oot_without_date_not_supported").format(idx + 1)
+                elif self.columns_for_online_api:
+                    msg = self.bundle.get("oot_with_online_sources_not_supported").format(idx + 1)
+                else:
+                    msg = None
+                if msg:
+                    if not silent:
+                        print(msg)
+                        self.logger.warning(msg)
+                    continue
+            validated_eval_set.append(validated_pair)
 
-        return validated_oot
+        return validated_eval_set
 
     def _validate_eval_set_pair(self, X: pd.DataFrame, eval_pair: Tuple) -> Tuple[pd.DataFrame, pd.Series]:
         if len(eval_pair) != 2:
@@ -3562,16 +3816,18 @@ if response.status_code == 200:
             raise ValidationError(self.bundle.get("unsupported_y_type_eval_set").format(type(eval_y)))
 
         eval_y_nunique = validated_eval_y.nunique()
-        if eval_y_nunique < 2:
+        is_oot = validated_eval_y.isna().all()
+        if not is_oot and eval_y_nunique < 2:
             raise ValidationError(self.bundle.get("y_is_constant_eval_set"))
 
-        if self.model_task_type == ModelTaskType.BINARY and eval_y_nunique != 2:
+        if not is_oot and self.model_task_type == ModelTaskType.BINARY and eval_y_nunique != 2:
             raise ValidationError(self.bundle.get("binary_target_eval_unique_count_not_2").format(eval_y_nunique))
 
-        # Check for duplicates between train and eval sets by comparing all values
-        train_eval_intersection = pd.merge(X, validated_eval_X, how="inner")
-        if len(train_eval_intersection) > 0:
-            raise ValidationError(self.bundle.get("eval_x_has_train_samples"))
+        if not is_oot:
+            # Check for duplicates between train and eval sets by comparing all values
+            train_eval_intersection = pd.merge(X, validated_eval_X, how="inner")
+            if len(train_eval_intersection) > 0:
+                raise ValidationError(self.bundle.get("eval_x_has_train_samples"))
 
         return validated_eval_X, validated_eval_y
 
@@ -3587,10 +3843,12 @@ if response.status_code == 200:
                 if isinstance(eval_set, tuple):
                     eval_set = [eval_set]
                 for eval in eval_set:
-                    if self.baseline_score_column not in eval[0].columns:
-                        raise ValidationError(self.bundle.get("baseline_score_column_not_exists"))
-                    if eval[0][self.baseline_score_column].isna().any():
-                        raise ValidationError(self.bundle.get("baseline_score_column_has_na"))
+                    is_oot = eval[1].isna().all()
+                    if not is_oot:
+                        if self.baseline_score_column not in eval[0].columns:
+                            raise ValidationError(self.bundle.get("baseline_score_column_not_exists"))
+                        if eval[0][self.baseline_score_column].isna().any():
+                            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]:
@@ -3660,7 +3918,6 @@ if response.status_code == 200:
         X: pd.DataFrame,
         y: Union[pd.Series, np.ndarray, list, None] = None,
         eval_set: Optional[List[tuple]] = None,
-        oot: Optional[pd.DataFrame] = None,
         exclude_features_sources: Optional[List[str]] = None,
         calculate_metrics: Optional[bool] = None,
         cv: Optional[Any] = None,
@@ -3729,8 +3986,6 @@ if response.status_code == 200:
                         self.logger.info(
                             f"First 10 rows of the eval_y_{idx} with shape {_num_samples(eval_y)}:\n{sample(eval_y)}"
                         )
-                if oot is not None:
-                    self.logger.info(f"First 10 rows of the oot with shape {oot.shape}:\n{sample(oot)}")
 
             do_without_pandas_limits(print_datasets_sample)
 
@@ -3767,7 +4022,7 @@ if response.status_code == 200:
         return df
 
     def _add_current_date_as_key(
-        self, df: pd.DataFrame, search_keys: Dict[str, SearchKey], logger: logging.Logger, bundle: ResourceBundle
+        self, df: pd.DataFrame, search_keys: Dict[str, SearchKey], bundle: ResourceBundle, silent: bool = False
     ) -> pd.DataFrame:
         if (
             set(search_keys.values()) == {SearchKey.PHONE}
@@ -3775,7 +4030,8 @@ if response.status_code == 200:
             or set(search_keys.values()) == {SearchKey.HEM}
             or set(search_keys.values()) == {SearchKey.COUNTRY, SearchKey.POSTAL_CODE}
         ):
-            self.__log_warning(bundle.get("current_date_added"))
+            if not silent:
+                self.__log_warning(bundle.get("current_date_added"))
             df[FeaturesEnricher.CURRENT_DATE] = datetime.date.today()
             search_keys[FeaturesEnricher.CURRENT_DATE] = SearchKey.DATE
             converter = DateTimeSearchKeyConverter(FeaturesEnricher.CURRENT_DATE, generate_cyclical_features=False)
@@ -4110,7 +4366,11 @@ if response.status_code == 200:
         return [f.name for f in features_meta if f.type == "categorical" and f.name not in (self.id_columns or [])]
 
     def __prepare_feature_importances(
-        self, trace_id: str, df: pd.DataFrame, updated_shaps: Optional[Dict[str, float]] = None, silent=False
+        self,
+        trace_id: str,
+        clients_features_df: pd.DataFrame,
+        updated_shaps: Optional[Dict[str, float]] = None,
+        silent=False,
     ):
         if self._search_task is None:
             raise NotFittedError(self.bundle.get("transform_unfitted_enricher"))
@@ -4123,11 +4383,12 @@ if response.status_code == 200:
         features_df = self._search_task.get_all_initial_raw_features(trace_id, metrics_calculation=True)
 
         # To be sure that names with hash suffixes
-        df = df.rename(columns=original_names_dict)
+        clients_features_df = clients_features_df.rename(columns=original_names_dict)
 
         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 = []
@@ -4136,10 +4397,10 @@ if response.status_code == 200:
         original_shaps = {original_names_dict.get(fm.name, fm.name): fm.shap_value for fm in features_meta}
 
         for feature_meta in features_meta:
-            if feature_meta.name in original_names_dict.keys():
-                feature_meta.name = original_names_dict[feature_meta.name]
+            original_name = original_names_dict.get(feature_meta.name, feature_meta.name)
+            feature_meta.name = original_name
 
-            is_client_feature = original_names_dict.get(feature_meta.name, feature_meta.name) in df.columns
+            is_client_feature = original_name in clients_features_df.columns
 
             # 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):
@@ -4156,12 +4417,21 @@ if response.status_code == 200:
 
         for feature_meta in features_meta:
             original_name = original_names_dict.get(feature_meta.name, feature_meta.name)
-            is_client_feature = original_name in df.columns
+            is_client_feature = original_name in clients_features_df.columns
 
             if not is_client_feature:
                 self.external_source_feature_names.append(original_name)
 
+            if self.psi_values is not None:
+                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)
@@ -4185,7 +4455,7 @@ if response.status_code == 200:
             self.feature_names_.append(feature_meta.name)
             self.feature_importances_.append(_round_shap_value(feature_meta.shap_value))
 
-            df_for_sample = features_df if feature_meta.name in features_df.columns else df
+            df_for_sample = features_df if feature_meta.name in features_df.columns else clients_features_df
             feature_info = FeatureInfo.from_metadata(feature_meta, df_for_sample, is_client_feature)
             features_info.append(feature_info.to_row(self.bundle))
             features_info_without_links.append(feature_info.to_row_without_links(self.bundle))
@@ -4193,6 +4463,8 @@ if response.status_code == 200:
 
         if len(features_info) > 0:
             self.features_info = pd.DataFrame(features_info)
+            if self.features_info[self.bundle.get("features_info_psi")].isna().all():
+                self.features_info.drop(columns=[self.bundle.get("features_info_psi")], inplace=True)
             self._features_info_without_links = pd.DataFrame(features_info_without_links)
             self._internal_features_info = pd.DataFrame(internal_features_info)
             if not silent:
@@ -4640,7 +4912,7 @@ if response.status_code == 200:
             print(msg)
 
     def _validate_PSI(self, df: pd.DataFrame):
-        if EVAL_SET_INDEX in df.columns and (df[EVAL_SET_INDEX] == 1).any():
+        if EVAL_SET_INDEX in df.columns:
             train = df.query(f"{EVAL_SET_INDEX} == 0")
             eval1 = df.query(f"{EVAL_SET_INDEX} == 1")
         else: