upgini 1.2.113a3974.dev2__py3-none-any.whl → 1.2.114__py3-none-any.whl

upgini/utils/psi.py

@@ -0,0 +1,300 @@
+import itertools
+import logging
+import operator
+from functools import reduce
+from typing import Callable, Dict, Optional
+
+import more_itertools
+import numpy as np
+import pandas as pd
+from pandas.api.types import is_numeric_dtype
+from pydantic import BaseModel
+
+from upgini.metadata import TARGET, ModelTaskType
+
+
+class StabilityParams(BaseModel):
+    threshold: float = 999
+    n_intervals: int = 12
+    min_intervals: int = 10
+    max_intervals: Optional[int] = None
+    min_values_in_interval: Optional[int] = None
+    n_bins: int = 10
+    min_values_in_bin: Optional[int] = None
+    cat_top_pct: float = 0.7
+    agg: str = "max"
+
+
+DEFAULT_TARGET_PARAMS = StabilityParams(
+    n_intervals=12,
+    min_intervals=10,
+    max_intervals=None,
+    min_values_in_interval=None,
+    n_bins=5,
+)
+
+DEFAULT_FEATURES_PARAMS = StabilityParams(
+    n_intervals=12,
+    min_intervals=10,
+    max_intervals=None,
+    min_values_in_interval=None,
+    n_bins=10,
+)
+
+
+def calculate_sparsity_psi(
+    df: pd.DataFrame,
+    cat_features: list[str],
+    date_column: str,
+    logger: logging.Logger,
+    model_task_type: ModelTaskType,
+    stability_agg_func: str | None = None,
+    psi_features_params: StabilityParams = DEFAULT_FEATURES_PARAMS,
+    psi_target_params: StabilityParams = DEFAULT_TARGET_PARAMS,
+) -> Dict[str, float]:
+    sparse_features = df.columns[df.isna().sum() > 0].to_list()
+    if len(sparse_features) > 0:
+        logger.info(f"Calculating sparsity stability for {len(sparse_features)} sparse features")
+        sparse_df = df[sparse_features].notna()
+        sparse_df[date_column] = df[date_column]
+        return calculate_features_psi(
+            sparse_df,
+            cat_features,
+            date_column,
+            logger,
+            model_task_type,
+            stability_agg_func,
+            psi_target_params,
+            psi_features_params,
+        )
+    return {}
+
+
+def calculate_features_psi(
+    df: pd.DataFrame,
+    cat_features: list[str],
+    date_column: str,
+    logger: logging.Logger,
+    model_task_type: ModelTaskType,
+    stability_agg_func: str | None = None,
+    psi_features_params: StabilityParams = DEFAULT_FEATURES_PARAMS,
+    psi_target_params: StabilityParams = DEFAULT_TARGET_PARAMS,
+) -> dict[str, float]:
+    empty_res = {col: 0.0 for col in df.columns if col not in [TARGET, date_column]}
+
+    if not is_numeric_dtype(df[date_column]):
+        df[date_column] = pd.to_datetime(df[date_column]).dt.floor("D").astype(np.int64) / 10**6
+
+    # Filter out rows with missing dates
+    df = df[df[date_column].notna()].copy()
+
+    n_months = pd.to_datetime(df[date_column], unit="ms").dt.month.nunique()
+
+    if TARGET in df.columns:
+        psi_target_params.n_intervals = min(
+            psi_target_params.max_intervals or np.inf, max(psi_target_params.min_intervals, n_months)
+        )
+        logger.info(f"Setting {psi_target_params.n_intervals} intervals for target PSI check")
+
+        logger.info(f"Calculating target PSI for {psi_target_params.n_intervals} intervals")
+        reference_mask, current_masks = _split_intervals(df, date_column, psi_target_params.n_intervals, logger)
+
+        if psi_target_params.min_values_in_interval is not None and any(
+            len(mask) < psi_target_params.min_values_in_interval
+            for mask in itertools.chain(current_masks, [reference_mask])
+        ):
+            logger.info(
+                f"Some intervals have less than {psi_target_params.min_values_in_interval} values. Skip PSI check"
+            )
+            return empty_res
+
+        target_agg_func = _get_agg_func(stability_agg_func or psi_target_params.agg)
+        logger.info(f"Calculating target PSI with agg function {target_agg_func}")
+        target_psi = _stability_agg(
+            [df[TARGET][cur] for cur in current_masks],
+            reference_data=df[TARGET][reference_mask],
+            is_numerical=model_task_type == ModelTaskType.REGRESSION,
+            min_values_in_bin=psi_target_params.min_values_in_bin,
+            n_bins=psi_target_params.n_bins,
+            cat_top_pct=psi_target_params.cat_top_pct,
+            agg_func=target_agg_func,
+        )
+        if target_psi is None or np.isnan(target_psi):
+            logger.info("Cannot determine target PSI. Skip feature PSI check")
+            return empty_res
+
+        if target_psi > psi_target_params.threshold:
+            logger.info(
+                f"Target PSI {target_psi} is more than threshold {psi_target_params.threshold}. Skip feature PSI check"
+            )
+            return empty_res
+
+    psi_features_params.n_intervals = min(
+        psi_features_params.max_intervals or np.inf, max(psi_features_params.min_intervals, n_months)
+    )
+    logger.info(f"Setting {psi_features_params.n_intervals} intervals for features PSI check")
+
+    logger.info(f"Calculating PSI for {len(df.columns)} features")
+    reference_mask, current_masks = _split_intervals(df, date_column, psi_features_params.n_intervals, logger)
+    features_agg_func = _get_agg_func(stability_agg_func or psi_features_params.agg)
+    logger.info(f"Calculating features PSI with agg function {features_agg_func}")
+    psi_values = [
+        _stability_agg(
+            [df[feature][cur] for cur in current_masks],
+            reference_data=df[feature][reference_mask],
+            is_numerical=feature not in cat_features,
+            min_values_in_bin=psi_features_params.min_values_in_bin,
+            n_bins=psi_features_params.n_bins,
+            cat_top_pct=psi_features_params.cat_top_pct,
+            agg_func=features_agg_func,
+        )
+        for feature in df.columns
+        if feature not in [TARGET, date_column]
+    ]
+    return {feature: psi for feature, psi in zip(df.columns, psi_values)}
+
+
+def _split_intervals(
+    df: pd.DataFrame, date_column: str, n_intervals: int, logger: logging.Logger
+) -> tuple[pd.Series, list[pd.Series]]:
+    date_series = df[date_column]
+
+    # Check if we have enough unique values for the requested number of intervals
+    unique_values = date_series.nunique()
+
+    # If we have fewer unique values than requested intervals, adjust n_intervals
+    if unique_values < n_intervals:
+        logger.warning(f"Date column '{date_column}' has only {unique_values} unique values")
+
+    time_intervals = pd.qcut(date_series, q=n_intervals, duplicates="drop")
+    interval_labels = time_intervals.unique()
+    reference_mask = time_intervals == interval_labels[0]
+    current_masks = [time_intervals == label for label in interval_labels[1:]]
+    return reference_mask, current_masks
+
+
+def _get_agg_func(agg: str):
+    np_agg = getattr(np, agg, None)
+    if np_agg is None and agg.startswith("q"):
+        q = int(agg[1:])
+        return lambda x: np.quantile(list(x), q / 100, method="higher")
+    return np_agg
+
+
+def _psi(reference_percent: np.ndarray, current_percent: np.ndarray) -> float:
+    return np.sum((reference_percent - current_percent) * np.log(reference_percent / current_percent))
+
+
+def _stability_agg(
+    current_data: list[pd.Series],
+    reference_data: pd.Series,
+    is_numerical: bool = True,
+    min_values_in_bin: int | None = None,
+    n_bins: int = 10,
+    cat_top_pct: float = 0.7,
+    agg_func: Callable = max,
+) -> float | None:
+    """Calculate the PSI
+    Args:
+        current_data: current data
+        reference_data: reference data
+        is_numerical: whether the feature is numerical
+        reference_ratio: ratio of current data to use as reference if reference_data is not provided
+        min_values_in_bin: minimum number of values in a bin to calculate PSI
+        n_bins: number of bins to use for numerical features
+    Returns:
+        psi_value: calculated PSI
+    """
+    reference, current = _get_binned_data(reference_data, current_data, is_numerical, n_bins, cat_top_pct)
+
+    if len(reference) == 0 or len(current) == 0:
+        return None
+
+    nonempty_current = [i for i, c in enumerate(current) if len(c) > 0]
+    current = [current[i] for i in nonempty_current]
+    current_data = [current_data[i] for i in nonempty_current]
+
+    if len(current) == 0:
+        return None
+
+    if min_values_in_bin is not None and (
+        np.array(reference).min() < min_values_in_bin or any(np.array(c).min() < min_values_in_bin for c in current)
+    ):
+        return None
+
+    reference = _fill_zeroes(reference / len(reference_data))
+    current = [_fill_zeroes(c / len(d)) for c, d in zip(current, current_data)]
+
+    psi_value = agg_func([_psi(reference, c) for c in current])
+
+    return float(psi_value)
+
+
+def _get_binned_data(
+    reference_data: pd.Series,
+    current_data: list[pd.Series],
+    is_numerical: bool,
+    n_bins: int,
+    cat_top_pct: float,
+):
+    """Split variable into n buckets based on reference quantiles
+    Args:
+        reference_data: reference data
+        current_data: current data
+        feature_type: feature type
+        n: number of quantiles
+    Returns:
+        reference_counts: number of records in each bucket for reference
+        current_counts: number of records in each bucket for current
+    """
+    n_vals = reference_data.nunique()
+
+    if is_numerical and n_vals > 20:
+        bins = _get_bin_edges(reference_data, n_bins)
+        reference_counts = np.histogram(reference_data, bins)[0]
+        current_counts = [np.histogram(d, bins)[0] for d in current_data]
+
+    else:
+        keys = _get_unique_not_nan_values_list_from_series([reference_data] + current_data)
+        ref_feature_dict = {**dict.fromkeys(keys, 0), **dict(reference_data.value_counts())}
+        current_feature_dict = [{**dict.fromkeys(keys, 0), **dict(d.value_counts())} for d in current_data]
+        key_dict = more_itertools.map_reduce(
+            itertools.chain(ref_feature_dict.items(), *(d.items() for d in current_feature_dict)),
+            keyfunc=operator.itemgetter(0),
+            valuefunc=operator.itemgetter(1),
+            reducefunc=sum,
+        )
+        key_dict = pd.Series(key_dict)
+        keys = key_dict.index[key_dict.rank(pct=True) >= cat_top_pct]
+        reference_counts = np.array([ref_feature_dict[key] for key in keys])
+        current_counts = [np.array([current_feature_dict[i][key] for key in keys]) for i in range(len(current_data))]
+
+    reference_counts = np.append(reference_counts, reference_data.isna().sum())
+    current_counts = [np.append(d, current_data[i].isna().sum()) for i, d in enumerate(current_counts)]
+
+    return reference_counts, current_counts
+
+
+def _fill_zeroes(percents: np.ndarray) -> np.ndarray:
+    eps = 0.0001
+    if (percents == 0).all():
+        np.place(percents, percents == 0, eps)
+    else:
+        min_value = min(percents[percents != 0])
+        if min_value <= eps:
+            np.place(percents, percents == 0, eps)
+        else:
+            np.place(percents, percents == 0, min_value / 10**6)
+    return percents
+
+
+def _get_bin_edges(data: pd.Series, n_bins: int) -> np.ndarray:
+    bins = np.nanquantile(data, np.linspace(0, 1, n_bins + 1))
+    bins[0] = -np.inf
+    bins[-1] = np.inf
+    return bins
+
+
+def _get_unique_not_nan_values_list_from_series(series: list[pd.Series]) -> list:
+    """Get unique values from current and reference series, drop NaNs"""
+    return list(reduce(set.union, (set(s.dropna().unique()) for s in series)))

upgini/utils/sample_utils.py

@@ -1,55 +1,30 @@
-from dataclasses import dataclass, field
 import logging
 import numbers
+from dataclasses import dataclass
 from typing import Callable, List, Optional
+
 import numpy as np
 import pandas as pd
 
-from upgini.metadata import SYSTEM_RECORD_ID, CVType, ModelTaskType
+from upgini.metadata import (
+    EVAL_SET_INDEX,
+    SYSTEM_RECORD_ID,
+    TARGET,
+    CVType,
+    ModelTaskType,
+)
 from upgini.resource_bundle import ResourceBundle, get_custom_bundle
+from upgini.utils.config import (
+    TS_DEFAULT_HIGH_FREQ_TRUNC_LENGTHS,
+    TS_DEFAULT_LOW_FREQ_TRUNC_LENGTHS,
+    TS_DEFAULT_TIME_UNIT_THRESHOLD,
+    TS_MIN_DIFFERENT_IDS_RATIO,
+    SampleConfig,
+)
 from upgini.utils.target_utils import balance_undersample
 from upgini.utils.ts_utils import get_most_frequent_time_unit, trunc_datetime
 
 
-TS_MIN_DIFFERENT_IDS_RATIO = 0.2
-TS_DEFAULT_HIGH_FREQ_TRUNC_LENGTHS = [pd.DateOffset(years=2, months=6), pd.DateOffset(years=2, days=7)]
-TS_DEFAULT_LOW_FREQ_TRUNC_LENGTHS = [pd.DateOffset(years=7), pd.DateOffset(years=5)]
-TS_DEFAULT_TIME_UNIT_THRESHOLD = pd.Timedelta(weeks=4)
-FIT_SAMPLE_ROWS_TS = 100_000
-
-BINARY_MIN_SAMPLE_THRESHOLD = 5_000
-MULTICLASS_MIN_SAMPLE_THRESHOLD = 25_000
-BINARY_BOOTSTRAP_LOOPS = 5
-MULTICLASS_BOOTSTRAP_LOOPS = 2
-
-FIT_SAMPLE_THRESHOLD = 100_000
-FIT_SAMPLE_ROWS = 100_000
-FIT_SAMPLE_ROWS_WITH_EVAL_SET = 100_000
-FIT_SAMPLE_THRESHOLD_WITH_EVAL_SET = 100_000
-
-
-@dataclass
-class SampleConfig:
-    force_sample_size: int = 7000
-    ts_min_different_ids_ratio: float = TS_MIN_DIFFERENT_IDS_RATIO
-    ts_default_high_freq_trunc_lengths: List[pd.DateOffset] = field(
-        default_factory=TS_DEFAULT_HIGH_FREQ_TRUNC_LENGTHS.copy
-    )
-    ts_default_low_freq_trunc_lengths: List[pd.DateOffset] = field(
-        default_factory=TS_DEFAULT_LOW_FREQ_TRUNC_LENGTHS.copy
-    )
-    ts_default_time_unit_threshold: pd.Timedelta = TS_DEFAULT_TIME_UNIT_THRESHOLD
-    binary_min_sample_threshold: int = BINARY_MIN_SAMPLE_THRESHOLD
-    multiclass_min_sample_threshold: int = MULTICLASS_MIN_SAMPLE_THRESHOLD
-    binary_bootstrap_loops: int = BINARY_BOOTSTRAP_LOOPS
-    multiclass_bootstrap_loops: int = MULTICLASS_BOOTSTRAP_LOOPS
-    fit_sample_threshold: int = FIT_SAMPLE_THRESHOLD
-    fit_sample_rows: int = FIT_SAMPLE_ROWS
-    fit_sample_rows_with_eval_set: int = FIT_SAMPLE_ROWS_WITH_EVAL_SET
-    fit_sample_threshold_with_eval_set: int = FIT_SAMPLE_THRESHOLD_WITH_EVAL_SET
-    fit_sample_rows_ts: int = FIT_SAMPLE_ROWS_TS
-
-
 @dataclass
 class SampleColumns:
     date: str
@@ -117,6 +92,22 @@ def sample(
             **kwargs,
         )
 
+    # separate OOT
+    oot_dfs = []
+    other_dfs = []
+    if EVAL_SET_INDEX in df.columns:
+        for eval_set_index in df[EVAL_SET_INDEX].unique():
+            eval_df = df[df[EVAL_SET_INDEX] == eval_set_index]
+            if TARGET in eval_df.columns and eval_df[TARGET].isna().all():
+                oot_dfs.append(eval_df)
+            else:
+                other_dfs.append(eval_df)
+    if len(oot_dfs) > 0:
+        oot_df = pd.concat(oot_dfs, ignore_index=False)
+        df = pd.concat(other_dfs, ignore_index=False)
+    else:
+        oot_df = None
+
     num_samples = _num_samples(df)
     if num_samples > fit_sample_threshold:
         logger.info(
@@ -126,6 +117,18 @@ def sample(
         df = df.sample(n=fit_sample_rows, random_state=random_state)
         logger.info(f"Shape after threshold resampling: {df.shape}")
 
+    if oot_df is not None:
+        num_samples_oot = _num_samples(oot_df)
+        if num_samples_oot > fit_sample_threshold:
+            logger.info(
+                f"OOT has size {num_samples_oot} more than threshold {fit_sample_threshold} "
+                f"and will be downsampled to {fit_sample_rows}"
+            )
+            oot_df = oot_df.sample(n=fit_sample_rows, random_state=random_state)
+        df = pd.concat([df, oot_df], ignore_index=False)
+
+    logger.info(f"Dataset size after downsampling: {len(df)}")
+
     return df
 
 
@@ -175,7 +178,7 @@ def sample_time_series_train_eval(
             )
         if logger is not None:
             logger.info(f"Eval set size: {len(eval_df)}")
-        df = pd.concat([train_df, eval_df])
+        df = pd.concat([train_df, eval_df], ignore_index=False)
 
     elif len(train_df) > max_rows:
         df = sample_time_series_trunc(

upgini/utils/target_utils.py

@@ -6,9 +6,14 @@ import pandas as pd
 from pandas.api.types import is_bool_dtype, is_datetime64_any_dtype, is_numeric_dtype
 
 from upgini.errors import ValidationError
-from upgini.metadata import SYSTEM_RECORD_ID, ModelTaskType
-from upgini.resource_bundle import ResourceBundle, get_custom_bundle, bundle
+from upgini.metadata import EVAL_SET_INDEX, SYSTEM_RECORD_ID, TARGET, ModelTaskType
+from upgini.resource_bundle import ResourceBundle, bundle, get_custom_bundle
 from upgini.sampler.random_under_sampler import RandomUnderSampler
+from upgini.utils.config import SampleConfig
+
+MAX_MULTICLASS_CLASS_COUNT = 100
+MIN_TARGET_CLASS_ROWS = 100
+IMBALANCE_THESHOLD = 0.6
 
 
 def prepare_target(y: Union[pd.Series, np.ndarray], target_type: ModelTaskType) -> Union[pd.Series, np.ndarray]:
@@ -106,6 +111,47 @@ def define_task(
     return task
 
 
+def is_imbalanced(
+    data: pd.DataFrame,
+    task_type: ModelTaskType,
+    sample_config: SampleConfig,
+    bundle: ResourceBundle,
+) -> bool:
+    if task_type is None or not task_type.is_classification():
+        return False
+
+    data = data.drop_duplicates(keep="first")
+    columns_without_target = [col for col in data.columns if col != TARGET]
+    data = data.drop_duplicates(subset=columns_without_target, keep=False)
+
+    if task_type == ModelTaskType.BINARY and len(data) <= sample_config.binary_min_sample_threshold:
+        return False
+
+    count = len(data)
+    target = data[TARGET]
+    target_classes_count = target.nunique()
+
+    if target_classes_count > MAX_MULTICLASS_CLASS_COUNT:
+        msg = bundle.get("dataset_to_many_multiclass_targets").format(target_classes_count, MAX_MULTICLASS_CLASS_COUNT)
+        raise ValidationError(msg)
+
+    vc = target.value_counts()
+    min_class_value = vc.index[len(vc) - 1]
+    min_class_count = vc[min_class_value]
+
+    if min_class_count < MIN_TARGET_CLASS_ROWS:
+        msg = bundle.get("dataset_rarest_class_less_min").format(
+            min_class_value, min_class_count, MIN_TARGET_CLASS_ROWS
+        )
+        raise ValidationError(msg)
+
+    min_class_percent = IMBALANCE_THESHOLD / target_classes_count
+    min_class_threshold = min_class_percent * count
+
+    # If min class count less than 30% for binary or (60 / classes_count)% for multiclass
+    return bool(min_class_count < min_class_threshold)
+
+
 def is_int_encoding(unique_values):
     return set(unique_values) == set(range(len(unique_values))) or set(unique_values) == set(
         range(1, len(unique_values) + 1)
@@ -132,6 +178,11 @@ def balance_undersample(
     if SYSTEM_RECORD_ID not in df.columns:
         raise Exception("System record id must be presented for undersampling")
 
+    # Rebalance and send to server only train data
+    # because eval set data will be sent separately in transform for metrics
+    if EVAL_SET_INDEX in df.columns:
+        df = df[df[EVAL_SET_INDEX] == 0]
+
     target = df[target_column].copy()
 
     vc = target.value_counts()

{upgini-1.2.113a3974.dev2.dist-info → upgini-1.2.114.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.113a3974.dev2
+Version: 1.2.114
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/
@@ -38,7 +38,7 @@ Requires-Dist: python-bidi==0.4.2
 Requires-Dist: python-dateutil>=2.8.0
 Requires-Dist: python-json-logger>=3.3.0
 Requires-Dist: requests>=2.8.0
-Requires-Dist: scikit-learn>=1.3.0
+Requires-Dist: scikit-learn<1.8.0,>=1.3.0
 Requires-Dist: scipy>=1.10.0
 Requires-Dist: shap>=0.44.0
 Requires-Dist: xhtml2pdf<0.3.0,>=0.2.11
@@ -270,9 +270,9 @@ from upgini.metadata import SearchKey
 enricher = FeaturesEnricher(
 	search_keys={
 		"subscription_activation_date": SearchKey.DATE,
-    		"country": SearchKey.COUNTRY,
-    		"zip_code": SearchKey.POSTAL_CODE,
-    		"hashed_email": SearchKey.HEM,
+		"country": SearchKey.COUNTRY,
+		"zip_code": SearchKey.POSTAL_CODE,
+		"hashed_email": SearchKey.HEM,
 		"last_visit_ip_address": SearchKey.IP,
 		"registered_with_phone": SearchKey.PHONE
 	})
@@ -358,9 +358,9 @@ from upgini.metadata import SearchKey
 enricher = FeaturesEnricher(
 	search_keys={
 		"subscription_activation_date": SearchKey.DATE,
-    		"country": SearchKey.COUNTRY,
-    		"zip_code": SearchKey.POSTAL_CODE,
-    		"hashed_email": SearchKey.HEM,
+		"country": SearchKey.COUNTRY,
+		"zip_code": SearchKey.POSTAL_CODE,
+		"hashed_email": SearchKey.HEM,
 		"last_visit_ip_address": SearchKey.IP,
 		"registered_with_phone": SearchKey.PHONE
 	}, 
@@ -381,7 +381,7 @@ from upgini.metadata import SearchKey
 enricher = FeaturesEnricher(
 	search_keys={
 		"subscription_activation_date": SearchKey.DATE,
-    		"zip_code": SearchKey.POSTAL_CODE,
+		"zip_code": SearchKey.POSTAL_CODE,
 	}, 
 	country_code = "US",
 	date_format = "%Y-%d-%m"
@@ -409,8 +409,8 @@ y = train_df["churn_flag"]
 enricher = FeaturesEnricher(
 	search_keys={
 		"subscription_activation_date": SearchKey.DATE,
-    		"country": SearchKey.COUNTRY,
-    		"zip_code": SearchKey.POSTAL_CODE
+		"country": SearchKey.COUNTRY,
+		"zip_code": SearchKey.POSTAL_CODE
 	})
 
 # everything is ready to fit! For 200к records fitting should take around 10 minutes,
@@ -464,8 +464,8 @@ And then, for `transform` in a production ML pipeline, you'll get enrichment wit
 enricher = FeaturesEnricher(
 	search_keys={
 		"subscription_activation_date": SearchKey.DATE,
-    		"country": SearchKey.COUNTRY,
-    		"zip_code": SearchKey.POSTAL_CODE,
+		"country": SearchKey.COUNTRY,
+		"zip_code": SearchKey.POSTAL_CODE,
 	},
 ) 
 ```
@@ -516,8 +516,8 @@ enricher = FeaturesEnricher(
 If you're working with multivariate time series, you should specify id columns of individual univariate series in `FeaturesEnricher`. For example, if you have a dataset predicting sales for different stores and products, you should specify store and product id columns as follows:
 ```python
 enricher = FeaturesEnricher(
-    search_keys={
-        "sales_date": SearchKey.DATE,
+		search_keys={
+				"sales_date": SearchKey.DATE,
     },
     id_columns=["store_id", "product_id"],
     cv=CVType.time_series
@@ -733,9 +733,52 @@ enricher.fit(
 )
 ```
 #### ⚠️ Requirements for out-of-time dataset  
-- Same data schema as for search initialization dataset  
+- Same data schema as for search initialization X dataset
 - Pandas dataframe representation
 
+There are 3 options to pass out-of-time without labels:
+```python
+enricher.fit(
+  train_ids_and_features,
+  train_label,
+  eval_set = [
+    (eval_ids_and_features_1,),  # Just tuple of 1 element
+    (eval_ids_and_features_2, None),  # None as labels
+    (eval_ids_and_features_3, [np.nan] * len(eval_ids_and_features_3)),  # List or Series of the same size as eval X
+  ]
+)
+```
+
+### Control feature stability with PSI parameters
+
+`FeaturesEnricher` supports Population Stability Index (PSI) calculation on eval_set to evaluate feature stability over time. You can control this behavior using stability parameters in `fit` and `fit_transform` methods:
+
+```python
+enricher = FeaturesEnricher(
+    search_keys={"registration_date": SearchKey.DATE}
+)
+
+# Control feature stability during fit
+enricher.fit(
+    X, y, 
+    stability_threshold=0.2,  # PSI threshold: features with PSI above this value will be dropped
+    stability_agg_func="max"  # Aggregation function for stability values: "max", "min", "mean"
+)
+
+# Same parameters work for fit_transform
+enriched_df = enricher.fit_transform(
+    X, y,
+    stability_threshold=0.1,   # Stricter threshold for more stable features
+    stability_agg_func="mean"  # Use mean aggregation instead of max
+)
+```
+
+**Stability parameters:**
+- `stability_threshold` (float, default=0.2): PSI threshold value. Features with PSI below this threshold will be excluded from the final feature set. Lower values mean stricter stability requirements.
+- `stability_agg_func` (str, default="max"): Function to aggregate PSI values across time intervals. Options: "max" (most conservative), "min" (least conservative), "mean" (balanced approach).
+
+**PSI (Population Stability Index)** measures how much feature distribution changes over time. Lower PSI values indicate more stable features, which are generally more reliable for production ML models.
+
 ### Use custom loss function in feature selection & metrics calculation
 
 `FeaturesEnricher` can be initialized with additional string parameter `loss`. 
@@ -756,20 +799,6 @@ enricher = FeaturesEnricher(
 enriched_dataframe.fit(X, y)
 ```
 
-### Return initial dataframe enriched with TOP external features by importance
-
-`transform` and `fit_transform` methods of `FeaturesEnricher` can be used with two additional parameters:
-- `importance_threshold`: float = 0 - only features with *importance >= threshold* will be added to the output dataframe
-- `max_features`: int  - only first TOP N features by importance will be returned, where *N = max_features*  
-
-And `keep_input=True` will keep all initial columns from search dataset X:  
-```python
-enricher = FeaturesEnricher(
-	search_keys={"subscription_activation_date": SearchKey.DATE}
-)
-enriched_dataframe.fit_transform(X, y, keep_input=True, max_features=2)
-```
-
 ### Exclude premium data sources from fit, transform and metrics calculation
 
 `fit`, `fit_transform`, `transform` and `calculate_metrics` methods of `FeaturesEnricher` can be used with parameter `exclude_features_sources` that allows to exclude Trial or Paid features from Premium data sources:
@@ -797,7 +826,7 @@ enricher = FeaturesEnricher(
 enricher.fit(X, y)
 ```
 
-## Turn off removing of target outliers
+### Turn off removing of target outliers
 Upgini detect rows with target outlier for regression tasks. By default such rows are dropped on metrics calculation. To turn off removing of target outlier rows use parameter `remove_outliers_calc_metrics=False` in fit, fit_transform or calculate_metrics methods:
 
 ```python
@@ -808,7 +837,7 @@ enricher = FeaturesEnricher(
 enricher.fit(X, y, remove_outliers_calc_metrics=False)
 ```
 
-## Turn off generating features on search keys
+### Turn off generating features on search keys
 Upgini tries to generate features on email, date and datetime search keys. By default this generation is enabled. To disable it use parameter `generate_search_key_features` of FeaturesEnricher constructor:
 
 ```python
@@ -816,6 +845,7 @@ enricher = FeaturesEnricher(
   search_keys={"date": SearchKey.DATE},
   generate_search_key_features=False,
 )
+```
 
 ## 🔑 Open up all capabilities of Upgini