upgini 1.1.222a1__py3-none-any.whl → 1.1.224__py3-none-any.whl

upgini/features_enricher.py

@@ -5,10 +5,13 @@ import logging
 import numbers
 import os
 import pickle
+import re
 import sys
 import tempfile
 import time
 import uuid
+from collections import namedtuple
+from functools import reduce
 from threading import Thread
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
@@ -175,6 +178,7 @@ class FeaturesEnricher(TransformerMixin):
         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,
         **kwargs,
     ):
@@ -279,6 +283,7 @@ class FeaturesEnricher(TransformerMixin):
 
         self.raise_validation_error = raise_validation_error
         self.exclude_columns = exclude_columns
+        self.baseline_score_column = baseline_score_column
 
     def _get_api_key(self):
         return self._api_key
@@ -290,6 +295,18 @@ class FeaturesEnricher(TransformerMixin):
 
     api_key = property(_get_api_key, _set_api_key)
 
+    @staticmethod
+    def _check_eval_set(eval_set, X):
+        checked_eval_set = []
+        if eval_set is not None and not isinstance(eval_set, list):
+            raise ValidationError(bundle.get("unsupported_type_eval_set").format(type(eval_set)))
+        for eval_pair in eval_set or []:
+            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 not is_frames_equal(X, eval_pair[0]):
+                checked_eval_set.append(eval_pair)
+        return checked_eval_set
+
     def fit(
         self,
         X: Union[pd.DataFrame, pd.Series, np.ndarray],
@@ -366,19 +383,13 @@ class FeaturesEnricher(TransformerMixin):
             try:
                 self.X = X
                 self.y = y
-                checked_eval_set = []
-                for eval_pair in eval_set or []:
-                    if len(eval_pair) != 2:
-                        raise ValidationError(bundle.get("eval_set_invalid_tuple_size").format(len(eval_pair)))
-                    if not is_frames_equal(X, eval_pair[0]):
-                        checked_eval_set.append(eval_pair)
-                self.eval_set = checked_eval_set
+                self.eval_set = self._check_eval_set(eval_set, X)
                 self.dump_input(trace_id, X, y, eval_set)
                 self.__inner_fit(
                     trace_id,
                     X,
                     y,
-                    checked_eval_set,
+                    self.eval_set,
                     progress_bar,
                     start_time=start_time,
                     exclude_features_sources=exclude_features_sources,
@@ -508,13 +519,7 @@ class FeaturesEnricher(TransformerMixin):
             try:
                 self.X = X
                 self.y = y
-                checked_eval_set = []
-                for eval_pair in eval_set or []:
-                    if len(eval_pair) != 2:
-                        raise ValidationError(bundle.get("eval_set_invalid_tuple_size").format(len(eval_pair)))
-                    if not is_frames_equal(X, eval_pair[0]):
-                        checked_eval_set.append(eval_pair)
-                self.eval_set = checked_eval_set
+                self.eval_set = self._check_eval_set(eval_set, X)
                 self.dump_input(trace_id, X, y, eval_set)
 
                 if _num_samples(drop_duplicates(X)) > Dataset.MAX_ROWS:
@@ -524,7 +529,7 @@ class FeaturesEnricher(TransformerMixin):
                     trace_id,
                     X,
                     y,
-                    checked_eval_set,
+                    self.eval_set,
                     progress_bar,
                     start_time=start_time,
                     exclude_features_sources=exclude_features_sources,
@@ -882,6 +887,7 @@ class FeaturesEnricher(TransformerMixin):
                     enriched_y_sorted,
                     fitting_eval_set_dict,
                     search_keys,
+                    groups,
                 ) = prepared_data
 
                 gc.collect()
@@ -906,7 +912,13 @@ class FeaturesEnricher(TransformerMixin):
                         ).get_cv()
 
                     wrapper = EstimatorWrapper.create(
-                        estimator, self.logger, model_task_type, _cv, fitting_enriched_X, scoring
+                        estimator,
+                        self.logger,
+                        model_task_type,
+                        _cv,
+                        fitting_enriched_X,
+                        scoring,
+                        groups=groups,
                     )
                     metric = wrapper.metric_name
                     multiplier = wrapper.multiplier
@@ -931,8 +943,11 @@ class FeaturesEnricher(TransformerMixin):
                             scoring,
                             cat_features,
                             add_params=custom_loss_add_params,
+                            groups=groups,
+                        )
+                        etalon_metric = baseline_estimator.cross_val_predict(
+                            fitting_X, y_sorted, self.baseline_score_column
                         )
-                        etalon_metric = baseline_estimator.cross_val_predict(fitting_X, y_sorted)
                         self.logger.info(f"Baseline {metric} on train client features: {etalon_metric}")
 
                     # 2 Fit and predict with KFold Catboost model on enriched tds
@@ -952,6 +967,7 @@ class FeaturesEnricher(TransformerMixin):
                             scoring,
                             cat_features,
                             add_params=custom_loss_add_params,
+                            groups=groups,
                         )
                         enriched_metric = enriched_estimator.cross_val_predict(fitting_enriched_X, enriched_y_sorted)
                         self.logger.info(f"Enriched {metric} on train combined features: {enriched_metric}")
@@ -999,7 +1015,9 @@ class FeaturesEnricher(TransformerMixin):
                                     f"Calculate baseline {metric} on eval set {idx + 1} "
                                     f"on client features: {eval_X_sorted.columns.to_list()}"
                                 )
-                                etalon_eval_metric = baseline_estimator.calculate_metric(eval_X_sorted, eval_y_sorted)
+                                etalon_eval_metric = baseline_estimator.calculate_metric(
+                                    eval_X_sorted, eval_y_sorted, self.baseline_score_column
+                                )
                                 self.logger.info(
                                     f"Baseline {metric} on eval set {idx + 1} client features: {etalon_eval_metric}"
                                 )
@@ -1172,12 +1190,7 @@ class FeaturesEnricher(TransformerMixin):
         if X is None:
             return True, self.X, self.y, self.eval_set
 
-        checked_eval_set = []
-        for eval_pair in eval_set or []:
-            if len(eval_pair) != 2:
-                raise ValidationError(bundle.get("eval_set_invalid_tuple_size").format(len(eval_pair)))
-            if not is_frames_equal(X, eval_pair[0]):
-                checked_eval_set.append(eval_pair)
+        checked_eval_set = self._check_eval_set(eval_set, X)
 
         if (
             X is self.X
@@ -1214,181 +1227,28 @@ class FeaturesEnricher(TransformerMixin):
         progress_bar: Optional[ProgressBar] = None,
         progress_callback: Optional[Callable[[SearchProgress], Any]] = None,
     ):
-        is_demo_dataset = hash_input(X, y, eval_set) in DEMO_DATASET_HASHES
         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
         validated_X = self._validate_X(X)
         validated_y = self._validate_y(validated_X, y)
+        validated_eval_set = (
+            [self._validate_eval_set_pair(validated_X, eval_set_pair) for eval_set_pair in eval_set]
+            if eval_set
+            else None
+        )
 
-        eval_set_sampled_dict = dict()
-
-        if self.__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")
-            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys = self.__cached_sampled_datasets
-            if exclude_features_sources:
-                enriched_X = drop_existing_columns(enriched_X, exclude_features_sources)
-        elif len(self.feature_importances_) == 0:
-            self.logger.info("No external features selected. So use only input datasets for metrics calculation")
-            X_sampled, search_keys = self._extend_x(validated_X, is_demo_dataset)
-            y_sampled = validated_y
-            enriched_X = X_sampled
-            if eval_set is not None:
-                for idx in range(len(eval_set)):
-                    eval_X_sampled, _ = self._extend_x(eval_set[idx][0], is_demo_dataset)
-                    eval_y_sampled = eval_set[idx][1]
-                    enriched_eval_X = eval_X_sampled
-                    eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
-            self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
-        elif not self.imbalanced and not exclude_features_sources and is_input_same_as_fit:
-            self.logger.info("Dataset is not imbalanced, so use enriched_X from fit")
-            search_keys = self.fit_search_keys
-
-            rows_to_drop = None
-            task_type = self.model_task_type or define_task(validated_y, self.logger, silent=True)
-            if task_type == ModelTaskType.REGRESSION:
-                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(
-                        self.df_with_original_index,
-                        target_outliers_df,
-                        left_on=SYSTEM_RECORD_ID,
-                        right_on=SYSTEM_RECORD_ID,
-                        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 "
-                    msg = bundle.get("target_outliers_warning").format(len(target_outliers_df), top_outliers, not_msg)
-                    print(msg)
-                    self.logger.warning(msg)
-
-            enriched_Xy, enriched_eval_sets = self.__enrich(
-                self.df_with_original_index,
-                self._search_task.get_all_initial_raw_features(trace_id, metrics_calculation=True),
-                rows_to_drop=rows_to_drop,
-            )
-
-            enriched_X = drop_existing_columns(enriched_Xy, TARGET)
-            x_columns = [
-                c for c in validated_X.columns.to_list() + self.fit_generated_features if c in enriched_X.columns
-            ]
-            X_sampled = enriched_Xy[x_columns].copy()
-            y_sampled = enriched_Xy[TARGET].copy()
-
-            self.logger.info(f"Shape of enriched_X: {enriched_X.shape}")
-            self.logger.info(f"Shape of X after sampling: {X_sampled.shape}")
-            self.logger.info(f"Shape of y after sampling: {len(y_sampled)}")
-
-            if eval_set is not None:
-                if len(enriched_eval_sets) != len(eval_set):
-                    raise ValidationError(
-                        bundle.get("metrics_eval_set_count_diff").format(len(enriched_eval_sets), len(eval_set))
-                    )
-
-                for idx in range(len(eval_set)):
-                    enriched_eval_X = drop_existing_columns(enriched_eval_sets[idx + 1], TARGET)
-                    eval_X_sampled = enriched_eval_sets[idx + 1][x_columns].copy()
-                    eval_y_sampled = enriched_eval_sets[idx + 1][TARGET].copy()
-                    eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
-
-            self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
-        else:
-            self.logger.info("Dataset is imbalanced or exclude_features_sources or X was passed. Run transform")
-            print(bundle.get("prepare_data_for_metrics"))
-            if eval_set is not None:
-                self.logger.info("Transform with eval_set")
-                # concatenate X and eval_set with eval_set_index
-                df_with_eval_set_index = validated_X.copy()
-                df_with_eval_set_index[TARGET] = validated_y
-                df_with_eval_set_index[EVAL_SET_INDEX] = 0
-                for idx, eval_pair in enumerate(eval_set):
-                    eval_x, eval_y = self._validate_eval_set_pair(validated_X, eval_pair)
-                    eval_df_with_index = eval_x.copy()
-                    eval_df_with_index[TARGET] = eval_y
-                    eval_df_with_index[EVAL_SET_INDEX] = idx + 1
-                    df_with_eval_set_index = pd.concat([df_with_eval_set_index, eval_df_with_index])
-
-                # downsample if need to eval_set threshold
-                num_samples = _num_samples(df_with_eval_set_index)
-                if num_samples > Dataset.FIT_SAMPLE_WITH_EVAL_SET_THRESHOLD:
-                    self.logger.info(f"Downsampling from {num_samples} to {Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS}")
-                    df_with_eval_set_index = df_with_eval_set_index.sample(
-                        n=Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS, random_state=self.random_state
-                    )
-
-                X_sampled = (
-                    df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == 0]
-                    .copy()
-                    .drop(columns=[EVAL_SET_INDEX, TARGET])
-                )
-                X_sampled, search_keys = self._extend_x(X_sampled, is_demo_dataset)
-                y_sampled = df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == 0].copy()[TARGET]
-                eval_set_sampled_dict = dict()
-                for idx in range(len(eval_set)):
-                    eval_x_sampled = (
-                        df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == (idx + 1)]
-                        .copy()
-                        .drop(columns=[EVAL_SET_INDEX, TARGET])
-                    )
-                    eval_x_sampled, _ = self._extend_x(eval_x_sampled, is_demo_dataset)
-                    eval_y_sampled = df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == (idx + 1)].copy()[
-                        TARGET
-                    ]
-                    eval_set_sampled_dict[idx] = (eval_x_sampled, eval_y_sampled)
-
-                df_with_eval_set_index.drop(columns=TARGET, inplace=True)
-
-                enriched = self.transform(
-                    df_with_eval_set_index,
-                    exclude_features_sources=exclude_features_sources,
-                    silent_mode=True,
-                    trace_id=trace_id,
-                    metrics_calculation=True,
-                    progress_bar=progress_bar,
-                    progress_callback=progress_callback,
-                )
-                if enriched is None:
-                    return None
-
-                enriched_X = enriched[enriched[EVAL_SET_INDEX] == 0].copy()
-                enriched_X.drop(columns=EVAL_SET_INDEX, inplace=True)
-
-                for idx in range(len(eval_set)):
-                    enriched_eval_x = enriched[enriched[EVAL_SET_INDEX] == (idx + 1)].copy()
-                    enriched_eval_x.drop(columns=EVAL_SET_INDEX, inplace=True)
-                    eval_x_sampled, eval_y_sampled = eval_set_sampled_dict[idx]
-                    eval_set_sampled_dict[idx] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
-            else:
-                self.logger.info("Transform without eval_set")
-                df = self.X.copy()
-
-                df[TARGET] = validated_y
-                num_samples = _num_samples(df)
-                if num_samples > Dataset.FIT_SAMPLE_THRESHOLD:
-                    self.logger.info(f"Downsampling from {num_samples} to {Dataset.FIT_SAMPLE_ROWS}")
-                    df = df.sample(n=Dataset.FIT_SAMPLE_ROWS, random_state=self.random_state)
-
-                X_sampled = df.copy().drop(columns=TARGET)
-                X_sampled, search_keys = self._extend_x(X_sampled, is_demo_dataset)
-                y_sampled = df.copy()[TARGET]
-
-                df.drop(columns=TARGET, inplace=True)
-
-                enriched_X = self.transform(
-                    df,
-                    exclude_features_sources=exclude_features_sources,
-                    silent_mode=True,
-                    trace_id=trace_id,
-                    metrics_calculation=True,
-                    progress_bar=progress_bar,
-                    progress_callback=progress_callback,
-                )
-                if enriched_X is None:
-                    return None
-
-            self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys = self._sample_data_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,
+        )
 
         excluding_search_keys = list(search_keys.keys())
         if search_keys_for_metrics is not None and len(search_keys_for_metrics) > 0:
@@ -1408,6 +1268,15 @@ class FeaturesEnricher(TransformerMixin):
         X_sorted, y_sorted = self._sort_by_keys(X_sampled, y_sampled, search_keys, self.cv)
         enriched_X_sorted, enriched_y_sorted = self._sort_by_keys(enriched_X, y_sampled, search_keys, self.cv)
 
+        group_columns = sorted(self._get_group_columns(search_keys))
+        groups = (
+            None
+            if not group_columns or self.cv != CVType.group_k_fold
+            else reduce(
+                lambda left, right: left + "_" + right, [enriched_X_sorted[c].astype(str) for c in group_columns]
+            ).factorize()[0]
+        )
+
         existing_filtered_enriched_features = [c for c in filtered_enriched_features if c in enriched_X_sorted.columns]
 
         fitting_X = X_sorted[client_features].copy()
@@ -1449,6 +1318,259 @@ class FeaturesEnricher(TransformerMixin):
             enriched_y_sorted,
             fitting_eval_set_dict,
             search_keys,
+            groups,
+        )
+
+    _SampledDataForMetrics = namedtuple(
+        "_SampledDataForMetrics", "X_sampled y_sampled enriched_X eval_set_sampled_dict search_keys"
+    )
+
+    def _sample_data_for_metrics(
+        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]],
+        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]],
+    ) -> _SampledDataForMetrics:
+        if self.__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(exclude_features_sources)
+        elif len(self.feature_importances_) == 0:
+            self.logger.info("No external features selected. So use only input datasets for metrics calculation")
+            return self.__sample_only_input(validated_X, validated_y, eval_set, is_demo_dataset)
+        elif not self.imbalanced and not exclude_features_sources and is_input_same_as_fit:
+            self.logger.info("Dataset is not imbalanced, so use enriched_X from fit")
+            return self.__sample_balanced(
+                validated_X, validated_y, eval_set, trace_id, remove_outliers_calc_metrics, is_demo_dataset
+            )
+        else:
+            self.logger.info("Dataset is imbalanced or exclude_features_sources or X was passed. Run transform")
+            print(bundle.get("prepare_data_for_metrics"))
+            return self.__sample_imbalanced(
+                validated_X,
+                validated_y,
+                eval_set,
+                is_demo_dataset,
+                exclude_features_sources,
+                trace_id,
+                progress_bar,
+                progress_callback,
+            )
+
+    def __get_sampled_cached_enriched(self, exclude_features_sources: Optional[List[str]]) -> _SampledDataForMetrics:
+        X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys = self.__cached_sampled_datasets
+        if exclude_features_sources:
+            enriched_X = drop_existing_columns(enriched_X, exclude_features_sources)
+
+        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+    def __sample_only_input(
+        self, validated_X: pd.DataFrame, validated_y: pd.Series, eval_set: Optional[List[tuple]], is_demo_dataset: bool
+    ) -> _SampledDataForMetrics:
+        eval_set_sampled_dict = dict()
+        X_sampled, search_keys = self._extend_x(validated_X, is_demo_dataset)
+        y_sampled = validated_y
+        enriched_X = X_sampled
+        if eval_set is not None:
+            for idx in range(len(eval_set)):
+                eval_X_sampled, _ = self._extend_x(eval_set[idx][0], is_demo_dataset)
+                eval_y_sampled = eval_set[idx][1]
+                enriched_eval_X = eval_X_sampled
+                eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
+        self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+    def __sample_balanced(
+        self,
+        validated_X: pd.DataFrame,
+        validated_y: pd.Series,
+        eval_set: Optional[List[tuple]],
+        trace_id: str,
+        remove_outliers_calc_metrics: Optional[bool],
+        is_demo_dataset: bool,
+    ) -> _SampledDataForMetrics:
+        eval_set_sampled_dict = dict()
+        search_keys = self.fit_search_keys
+
+        rows_to_drop = None
+        task_type = self.model_task_type or define_task(validated_y, self.logger, silent=True)
+        if task_type == ModelTaskType.REGRESSION:
+            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(
+                    self.df_with_original_index,
+                    target_outliers_df,
+                    left_on=SYSTEM_RECORD_ID,
+                    right_on=SYSTEM_RECORD_ID,
+                    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 "
+                msg = bundle.get("target_outliers_warning").format(len(target_outliers_df), top_outliers, not_msg)
+                print(msg)
+                self.logger.warning(msg)
+
+        # index in each dataset (X, eval set) may be reordered and non unique, but index in validated datasets
+        # can differs from it
+        enriched_Xy, enriched_eval_sets = self.__enrich(
+            self.df_with_original_index,
+            self._search_task.get_all_initial_raw_features(trace_id, metrics_calculation=True),
+            rows_to_drop=rows_to_drop,
+        )
+
+        enriched_X = drop_existing_columns(enriched_Xy, TARGET)
+        X_sampled, search_keys = self._extend_x(validated_X, is_demo_dataset)
+        y_sampled = enriched_Xy[TARGET].copy()
+
+        self.logger.info(f"Shape of enriched_X: {enriched_X.shape}")
+        self.logger.info(f"Shape of X after sampling: {X_sampled.shape}")
+        self.logger.info(f"Shape of y after sampling: {len(y_sampled)}")
+
+        if eval_set is not None:
+            if len(enriched_eval_sets) != len(eval_set):
+                raise ValidationError(
+                    bundle.get("metrics_eval_set_count_diff").format(len(enriched_eval_sets), len(eval_set))
+                )
+
+            for idx in range(len(eval_set)):
+                enriched_eval_X = drop_existing_columns(enriched_eval_sets[idx + 1], TARGET)
+                eval_X_sampled, _ = self._extend_x(eval_set[idx][0], is_demo_dataset)
+                eval_y_sampled = eval_set[idx][1].copy()
+                eval_set_sampled_dict[idx] = (eval_X_sampled, enriched_eval_X, eval_y_sampled)
+
+        self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+    def __sample_imbalanced(
+        self,
+        validated_X: pd.DataFrame,
+        validated_y: pd.Series,
+        eval_set: Optional[List[tuple]],
+        is_demo_dataset: bool,
+        exclude_features_sources: Optional[List[str]],
+        trace_id: str,
+        progress_bar: Optional[ProgressBar],
+        progress_callback: Optional[Callable[[SearchProgress], Any]],
+    ) -> _SampledDataForMetrics:
+        eval_set_sampled_dict = dict()
+        if eval_set is not None:
+            self.logger.info("Transform with eval_set")
+            # concatenate X and eval_set with eval_set_index
+            df_with_eval_set_index = validated_X.copy()
+            df_with_eval_set_index[TARGET] = validated_y
+            df_with_eval_set_index[EVAL_SET_INDEX] = 0
+            for idx, eval_pair in enumerate(eval_set):
+                eval_x, eval_y = self._validate_eval_set_pair(validated_X, eval_pair)
+                eval_df_with_index = eval_x.copy()
+                eval_df_with_index[TARGET] = eval_y
+                eval_df_with_index[EVAL_SET_INDEX] = idx + 1
+                df_with_eval_set_index = pd.concat([df_with_eval_set_index, eval_df_with_index])
+
+            # downsample if need to eval_set threshold
+            num_samples = _num_samples(df_with_eval_set_index)
+            if num_samples > Dataset.FIT_SAMPLE_WITH_EVAL_SET_THRESHOLD:
+                self.logger.info(f"Downsampling from {num_samples} to {Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS}")
+                df_with_eval_set_index = df_with_eval_set_index.sample(
+                    n=Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS, random_state=self.random_state
+                )
+
+            X_sampled = (
+                df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == 0]
+                .copy()
+                .drop(columns=[EVAL_SET_INDEX, TARGET])
+            )
+            X_sampled, search_keys = self._extend_x(X_sampled, is_demo_dataset)
+            y_sampled = df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == 0].copy()[TARGET]
+            eval_set_sampled_dict = dict()
+            for idx in range(len(eval_set)):
+                eval_x_sampled = (
+                    df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == (idx + 1)]
+                    .copy()
+                    .drop(columns=[EVAL_SET_INDEX, TARGET])
+                )
+                eval_x_sampled, _ = self._extend_x(eval_x_sampled, is_demo_dataset)
+                eval_y_sampled = df_with_eval_set_index[df_with_eval_set_index[EVAL_SET_INDEX] == (idx + 1)].copy()[
+                    TARGET
+                ]
+                eval_set_sampled_dict[idx] = (eval_x_sampled, eval_y_sampled)
+
+            df_with_eval_set_index.drop(columns=TARGET, inplace=True)
+
+            enriched = self.transform(
+                df_with_eval_set_index,
+                exclude_features_sources=exclude_features_sources,
+                silent_mode=True,
+                trace_id=trace_id,
+                metrics_calculation=True,
+                progress_bar=progress_bar,
+                progress_callback=progress_callback,
+            )
+            if enriched is None:
+                return None
+
+            enriched_X = enriched[enriched[EVAL_SET_INDEX] == 0].copy()
+            enriched_X.drop(columns=EVAL_SET_INDEX, inplace=True)
+
+            for idx in range(len(eval_set)):
+                enriched_eval_x = enriched[enriched[EVAL_SET_INDEX] == (idx + 1)].copy()
+                enriched_eval_x.drop(columns=EVAL_SET_INDEX, inplace=True)
+                eval_x_sampled, eval_y_sampled = eval_set_sampled_dict[idx]
+                eval_set_sampled_dict[idx] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
+        else:
+            self.logger.info("Transform without eval_set")
+            df = self.X.copy()
+
+            df[TARGET] = validated_y
+            num_samples = _num_samples(df)
+            if num_samples > Dataset.FIT_SAMPLE_THRESHOLD:
+                self.logger.info(f"Downsampling from {num_samples} to {Dataset.FIT_SAMPLE_ROWS}")
+                df = df.sample(n=Dataset.FIT_SAMPLE_ROWS, random_state=self.random_state)
+
+            X_sampled = df.copy().drop(columns=TARGET)
+            X_sampled, search_keys = self._extend_x(X_sampled, is_demo_dataset)
+            y_sampled = df.copy()[TARGET]
+
+            df.drop(columns=TARGET, inplace=True)
+
+            enriched_X = self.transform(
+                df,
+                exclude_features_sources=exclude_features_sources,
+                silent_mode=True,
+                trace_id=trace_id,
+                metrics_calculation=True,
+                progress_bar=progress_bar,
+                progress_callback=progress_callback,
+            )
+            if enriched_X is None:
+                return None
+
+        self.__cached_sampled_datasets = (X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+        return self.__mk_sampled_data_tuple(X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys)
+
+    def __mk_sampled_data_tuple(
+        self,
+        X_sampled: pd.DataFrame,
+        y_sampled: pd.Series,
+        enriched_X: pd.DataFrame,
+        eval_set_sampled_dict: Dict,
+        search_keys: Dict,
+    ):
+        search_keys = {k: v for k, v in search_keys.items() if k in X_sampled.columns.to_list()}
+        return FeaturesEnricher._SampledDataForMetrics(
+            X_sampled, y_sampled, enriched_X, eval_set_sampled_dict, search_keys
         )
 
     def get_search_id(self) -> Optional[str]:
@@ -1868,20 +1990,8 @@ class FeaturesEnricher(TransformerMixin):
 
         df = self.__add_country_code(df, self.fit_search_keys)
 
-        # Check Multivariate time series
         date_column = self._get_date_column(self.fit_search_keys)
-        if (
-            self.cv is None
-            and date_column
-            and model_task_type == ModelTaskType.REGRESSION
-            and len({SearchKey.PHONE, SearchKey.EMAIL, SearchKey.HEM}.intersection(self.fit_search_keys.keys())) == 0
-            and is_blocked_time_series(df, date_column, list(self.fit_search_keys.keys()) + [TARGET])
-        ):
-            msg = bundle.get("multivariate_timeseries_detected")
-            print(msg)
-            self.logger.warning(msg)
-            self.cv = CVType.blocked_time_series
-            self.runtime_parameters.properties["cv_type"] = self.cv.name
+        self.__adjust_cv(df, date_column, model_task_type)
 
         self.fit_generated_features = []
 
@@ -2046,8 +2156,9 @@ class FeaturesEnricher(TransformerMixin):
 
         self.__show_selected_features(self.fit_search_keys)
 
-        if not self.warning_counter.has_warnings():
-            self.__display_support_link(bundle.get("all_ok_community_invite"))
+        autofe_description = self.get_autofe_features_description()
+        if autofe_description is not None:
+            display_html_dataframe(autofe_description, autofe_description, "*Description of AutoFE feature names")
 
         if self._has_paid_features(exclude_features_sources):
             if calculate_metrics is not None and calculate_metrics:
@@ -2089,6 +2200,35 @@ class FeaturesEnricher(TransformerMixin):
 
         self.__show_report_button()
 
+        if not self.warning_counter.has_warnings():
+            self.__display_support_link(bundle.get("all_ok_community_invite"))
+
+    def __adjust_cv(self, df: pd.DataFrame, date_column: pd.Series, model_task_type: ModelTaskType):
+        # Check Multivariate time series
+        if (
+            self.cv is None
+            and date_column
+            and model_task_type == ModelTaskType.REGRESSION
+            and len({SearchKey.PHONE, SearchKey.EMAIL, SearchKey.HEM}.intersection(self.fit_search_keys.keys())) == 0
+            and is_blocked_time_series(df, date_column, list(self.fit_search_keys.keys()) + [TARGET])
+        ):
+            msg = bundle.get("multivariate_timeseries_detected")
+            self.__override_cv(CVType.blocked_time_series, msg, print_warning=False)
+        elif (
+            (self.cv is None or self.cv == CVType.k_fold)
+            and model_task_type != ModelTaskType.REGRESSION
+            and self._get_group_columns(self.fit_search_keys)
+        ):
+            msg = bundle.get("group_k_fold_in_classification")
+            self.__override_cv(CVType.group_k_fold, msg, print_warning=self.cv is not None)
+
+    def __override_cv(self, cv: CVType, msg: str, print_warning: bool = True):
+        if print_warning:
+            print(msg)
+        self.logger.warning(msg)
+        self.cv = cv
+        self.runtime_parameters.properties["cv_type"] = self.cv.name
+
     def get_columns_by_search_keys(self, keys: List[str]):
         if "HEM" in keys:
             keys.append("EMAIL")
@@ -2384,6 +2524,10 @@ class FeaturesEnricher(TransformerMixin):
             if t in [SearchKey.DATE, SearchKey.DATETIME]:
                 return col
 
+    @staticmethod
+    def _get_group_columns(search_keys: Dict[str, SearchKey]) -> List[str]:
+        return [col for col, t in search_keys.items() if t not in [SearchKey.DATE, SearchKey.DATETIME]]
+
     @staticmethod
     def __get_email_column(search_keys: Dict[str, SearchKey]) -> Optional[str]:
         for col, t in search_keys.items():
@@ -2400,8 +2544,10 @@ class FeaturesEnricher(TransformerMixin):
         self, df: pd.DataFrame, meaning_types: Dict[str, FileColumnMeaningType], search_keys: Dict[str, SearchKey]
     ) -> pd.DataFrame:
         # save original order or rows
-        df = df.reset_index(drop=True).reset_index()
-        df = df.rename(columns={DEFAULT_INDEX: ORIGINAL_INDEX})
+        original_index_name = df.index.name
+        index_name = df.index.name or DEFAULT_INDEX
+        df = df.reset_index().reset_index(drop=True)
+        df = df.rename(columns={index_name: ORIGINAL_INDEX})
 
         # order by date and idempotent order by other keys
         if self.cv not in [CVType.time_series, CVType.blocked_time_series]:
@@ -2432,7 +2578,8 @@ class FeaturesEnricher(TransformerMixin):
 
         # return original order
         df = df.set_index(ORIGINAL_INDEX)
-        df = df.sort_index()
+        df.index.name = original_index_name
+        # df = df.sort_index()
 
         meaning_types[SYSTEM_RECORD_ID] = FileColumnMeaningType.SYSTEM_RECORD_ID
         return df
@@ -2493,6 +2640,9 @@ class FeaturesEnricher(TransformerMixin):
             self.logger.warning(f"X contain columns with same name as returned from backend: {dup_features}")
             raise ValidationError(bundle.get("returned_features_same_as_passed").format(dup_features))
 
+        # index overrites from result_features
+        original_index_name = df_with_original_index.index.name
+        df_with_original_index = df_with_original_index.reset_index()
         result_features = pd.merge(
             df_with_original_index,
             result_features,
@@ -2500,6 +2650,8 @@ class FeaturesEnricher(TransformerMixin):
             right_on=SYSTEM_RECORD_ID,
             how="left" if is_transform else "inner",
         )
+        result_features = result_features.set_index(original_index_name or DEFAULT_INDEX)
+        result_features.index.name = original_index_name
 
         if rows_to_drop is not None:
             print(f"Before dropping target outliers size: {len(result_features)}")
@@ -2687,6 +2839,52 @@ class FeaturesEnricher(TransformerMixin):
         else:
             self.logger.warning("Empty features info")
 
+    def get_autofe_features_description(self):
+        try:
+            autofe_meta = self._search_task.get_autofe_metadata()
+            if autofe_meta is None:
+                return None
+            features_meta = self._search_task.get_all_features_metadata_v2()
+
+            def get_feature_by_display_index(idx):
+                for m in features_meta:
+                    if m.name.endswith(str(idx)):
+                        return m
+
+            descriptions = []
+            for m in autofe_meta:
+                description = dict()
+
+                feature_meta = get_feature_by_display_index(m.display_index)
+                if feature_meta is None:
+                    self.logger.warning(f"Feature meta for display index {m.display_index} not found")
+                    continue
+                description["Sources"] = feature_meta.data_source.replace("AutoFE: features from ", "")
+                description["Feature name"] = feature_meta.name
+
+                feature_idx = 1
+                for bc in m.base_columns:
+                    description[f"Feature {feature_idx}"] = bc.hashed_name
+                    feature_idx += 1
+
+                match = re.match(f"f_autofe_(.+)_{m.display_index}", feature_meta.name)
+                if match is None:
+                    self.logger.warning(f"Failed to infer autofe function from name {feature_meta.name}")
+                else:
+                    description["Function"] = match.group(1)
+
+                descriptions.append(description)
+
+            if len(descriptions) == 0:
+                return None
+
+            descriptions_df = pd.DataFrame(descriptions)
+            descriptions_df.fillna("", inplace=True)
+            return descriptions_df
+        except Exception:
+            self.logger.exception("Failed to generate AutoFE features description")
+            return None
+
     @staticmethod
     def _group_relevant_data_sources(df: pd.DataFrame) -> pd.DataFrame:
         return (
@@ -2889,8 +3087,10 @@ class FeaturesEnricher(TransformerMixin):
                 relevant_features_df=self._features_info_without_links,
                 relevant_datasources_df=self.relevant_data_sources,
                 metrics_df=self.metrics,
+                autofe_descriptions_df=self.get_autofe_features_description(),
                 search_id=self._search_task.search_task_id,
                 email=get_rest_client(self.endpoint, self.api_key).get_current_email(),
+                search_keys=[str(sk) for sk in self.search_keys.values()],
             )
         except Exception:
             pass
@@ -3014,19 +3214,15 @@ class FeaturesEnricher(TransformerMixin):
     def __display_support_link(self, link_text: Optional[str] = None):
         support_link = bundle.get("support_link")
         link_text = link_text or bundle.get("support_text")
-        # badge = bundle.get("slack_community_bage")
-        # alt = bundle.get("slack_community_alt")
         try:
             from IPython.display import HTML, display
 
             _ = get_ipython()  # type: ignore
             self.logger.warning(link_text)
-            print(link_text)
             display(
                 HTML(
-                    f"""<a href='{support_link}' target='_blank' rel='noopener noreferrer'>
-                    Support</a>"""
-                    # <img alt='{alt}' src='{badge}'></a>
+                    f"""{link_text} <a href='{support_link}' target='_blank' rel='noopener noreferrer'>
+                    here</a>"""
                 )
             )
         except (ImportError, NameError):