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

upgini/metadata.py

@@ -285,6 +285,7 @@ class FeaturesMetadataV2(BaseModel):
     doc_link: Optional[str] = None
     update_frequency: Optional[str] = None
     from_online_api: Optional[bool] = None
+    psi_value: Optional[float] = None
 
 
 class HitRateMetrics(BaseModel):

upgini/metrics.py

@@ -816,7 +816,8 @@ class CatBoostWrapper(EstimatorWrapper):
                 else:
                     encoded = cat_encoder.transform(x[self.cat_features])
                     cat_features = encoded.columns.to_list()
-                x.loc[:, self.cat_features] = encoded
+                x.drop(columns=encoded.columns, inplace=True, errors="ignore")
+                x[encoded.columns] = encoded
             else:
                 cat_features = self.cat_features
 
@@ -1175,7 +1176,10 @@ def _ext_mean_squared_log_error(y_true, y_pred, *, sample_weight=None, multioutp
     >>> mean_squared_log_error(y_true, y_pred, multioutput=[0.3, 0.7])
     0.060...
     """
-    _, y_true, y_pred, multioutput = _check_reg_targets(y_true, y_pred, multioutput)
+    try:
+        _, y_true, y_pred, multioutput = _check_reg_targets(y_true, y_pred, multioutput)
+    except TypeError:
+        _, y_true, y_pred, sample_weight, multioutput = _check_reg_targets(y_true, y_pred, sample_weight, multioutput)
     check_consistent_length(y_true, y_pred, sample_weight)
 
     if (y_true < 0).any():

upgini/resource_bundle/strings.properties

@@ -123,7 +123,7 @@ train_unstable_target=Your training sample contains an unstable target event, PS
 eval_unstable_target=Your training and evaluation samples have a difference in target distribution. PSI = {}. The results will be unstable. It is recommended to redesign the training and evaluation samples
     # eval set validation
 unsupported_type_eval_set=Unsupported type of eval_set: {}. It should be list of tuples with two elements: X and y
-eval_set_invalid_tuple_size=eval_set contains a tuple of size {}. It should contain only pairs of X and y
+eval_set_invalid_tuple_size=eval_set contains a tuple of size {}. It should contain only pairs of X and y or X only
 unsupported_x_type_eval_set=Unsupported type of X in eval_set: {}. Use pandas.DataFrame, pandas.Series or numpy.ndarray or list.
 eval_x_and_x_diff_shape=The column set in eval_set are differ from the column set in X
 unsupported_y_type_eval_set=Unsupported type of y in eval_set: {}. Use pandas.Series, numpy.ndarray or list
@@ -139,6 +139,8 @@ eval_x_is_empty=X in eval_set is empty.
 eval_y_is_empty=y in eval_set is empty.
 x_and_eval_x_diff_types=X and eval_set X has different types: {} and {}
 eval_x_has_train_samples=Eval set X has rows that are present in train set X
+oot_without_date_not_supported=Eval set {} provided as OOT but date column is missing. It will be ignored for stability check
+oot_with_online_sources_not_supported=Eval set {} provided as OOT and also provided columns for online API. It will be ignored for stability check
 
 baseline_score_column_not_exists=baseline_score_column {} doesn't exist in input dataframe
 baseline_score_column_has_na=baseline_score_column contains NaN. Clear it and and retry
@@ -163,6 +165,7 @@ dataset_too_many_rows_registered=X rows limit for transform is {}. Please sample
 dataset_empty_column_names=Some column names are empty. Add names please
 dataset_full_duplicates={:.5f}% of the rows are fully duplicated
 dataset_diff_target_duplicates={:.4f}% of rows ({}) in X and eval_set are duplicates with different y values. These rows will be deleted as incorrect\nSample of incorrect row indexes: {}
+dataset_diff_target_duplicates_oot={:.4f}% of rows ({}) in OOT eval_set are duplicates with train or another eval_set. These rows will be deleted from OOT\nSample of incorrect row indexes: {}
 dataset_train_diff_target_duplicates_fintech={:.4f}% of rows ({}) in X are duplicates, not taking into consideration dates, IP addresses and features from the training set, but have different y values. These rows have been removed to optimize search results.\nRemoved row indexes: {}
 dataset_eval_diff_target_duplicates_fintech={:.4f}% of rows ({}) in eval{} X are duplicates, not taking into consideration dates, IP addresses and features from the training set, but have different y values. These rows have been removed to optimize search results.\nRemoved row indexes: {}
 dataset_drop_old_dates=We don't have data before '2000-01-01' and removed all earlier records from the search dataset
@@ -183,6 +186,7 @@ dataset_invalid_column_type=Unsupported data type of column {}: {}
 dataset_invalid_filter=Unknown field in filter_features. Should be {'min_importance', 'max_psi', 'max_count', 'selected_features'}.
 dataset_too_big_file=Too big size of dataframe X for processing. Please reduce number of rows or columns
 dataset_transform_diff_fit=You try to enrich dataset that column names are different from the train dataset column names that you used on the fit stage. Please make the column names the same as in the train dataset and restart.
+oot_eval_set_too_small_after_dedup=OOT eval set {} has less than 1000 rows after deduplication. It will be ignored for stability check
 binary_small_dataset=The least populated class in Target contains less than 1000 rows.\nSmall numbers of observations may negatively affect the number of selected features and quality of your ML model.\nUpgini recommends you increase the number of observations in the least populated class.\n
 all_search_keys_invalid=All search keys are invalid
 all_emails_invalid=All values in column {} are invalid emails    # Metrics validation
@@ -255,6 +259,7 @@ features_info_provider=Provider
 features_info_source=Source
 features_info_name=Feature name
 features_info_shap=SHAP value
+features_info_psi=PSI value
 features_info_hitrate=Coverage %
 features_info_type=Type
 # Deprecated

upgini/sampler/base.py

@@ -1,6 +1,7 @@
 """
 Base class for the under-sampling method.
 """
+
 # Authors: Guillaume Lemaitre <g.lemaitre58@gmail.com>
 # License: MIT
 
@@ -12,6 +13,7 @@ import numpy as np
 from sklearn.base import BaseEstimator
 from sklearn.preprocessing import label_binarize
 from sklearn.utils.multiclass import check_classification_targets
+from sklearn.utils.validation import check_X_y
 
 from .utils import ArraysTransformer, check_sampling_strategy, check_target_type
 
@@ -125,7 +127,7 @@ class BaseSampler(SamplerMixin):
         if accept_sparse is None:
             accept_sparse = ["csr", "csc"]
         y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
-        X, y = self._validate_data(X, y, reset=True, accept_sparse=accept_sparse)
+        X, y = check_X_y(X, y, accept_sparse=accept_sparse, dtype=None, ensure_all_finite=False)
         return X, y, binarize_y
 
     def _more_tags(self):

upgini/sampler/random_under_sampler.py

@@ -80,14 +80,24 @@ RandomUnderSampler # doctest: +NORMALIZE_WHITESPACE
 
     def _check_X_y(self, X, y):
         y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
-        X, y = self._validate_data(
-            X,
-            y,
-            reset=True,
-            accept_sparse=["csr", "csc"],
-            dtype=None,
-            force_all_finite=False,
-        )
+        try:
+            X, y = self._validate_data(
+                X,
+                y,
+                reset=True,
+                accept_sparse=["csr", "csc"],
+                dtype=None,
+                force_all_finite=False,
+            )
+        except AttributeError:
+            from sklearn.utils.validation import check_X_y
+            X, y = check_X_y(
+                X,
+                y,
+                accept_sparse=["csr", "csc"],
+                dtype=None,
+                ensure_all_finite=False,
+            )
         return X, y, binarize_y
 
     def _fit_resample(self, X, y):

upgini/utils/deduplicate_utils.py

@@ -134,8 +134,13 @@ def remove_fintech_duplicates(
     logger.info(f"Train dataset shape after clean fintech duplicates: {train_df.shape}")
 
     # Process each eval_set part separately
+    oot_eval_dfs = []
     new_eval_dfs = []
     for i, eval_df in enumerate(eval_dfs, 1):
+        # Skip OOT
+        if eval_df[TARGET].isna().all():
+            oot_eval_dfs.append(eval_df)
+            continue
         logger.info(f"Eval {i} dataset shape before clean fintech duplicates: {eval_df.shape}")
         cleaned_eval_df, eval_warning = process_df(eval_df, i)
         if eval_warning:
@@ -145,8 +150,8 @@ def remove_fintech_duplicates(
 
     # Combine the processed train and eval parts back into one dataset
     logger.info(f"Dataset shape before clean fintech duplicates: {df.shape}")
-    if new_eval_dfs:
-        df = pd.concat([train_df] + new_eval_dfs)
+    if new_eval_dfs or oot_eval_dfs:
+        df = pd.concat([train_df] + new_eval_dfs + oot_eval_dfs, ignore_index=False)
     else:
         df = train_df
     logger.info(f"Dataset shape after clean fintech duplicates: {df.shape}")
@@ -190,16 +195,59 @@ def clean_full_duplicates(
     msg = None
     if TARGET in df.columns:
         unique_columns.remove(TARGET)
-        marked_duplicates = df.duplicated(subset=unique_columns, keep=False)
+
+        # Separate rows to exclude from deduplication:
+        # for each eval_set_index != 0 check separately, all TARGET values are NaN
+        df_for_dedup = df
+        oot_df = None
+
+        if EVAL_SET_INDEX in df.columns:
+            oot_eval_dfs = []
+            other_dfs = []
+            for eval_idx in df[EVAL_SET_INDEX].unique():
+                eval_subset = df[df[EVAL_SET_INDEX] == eval_idx]
+                # Check that all TARGET values for this specific eval_set_index are NaN
+                if eval_idx != 0 and eval_subset[TARGET].isna().all():
+                    oot_eval_dfs.append(eval_subset)
+                    logger.info(
+                        f"Excluded {len(eval_subset)} rows from deduplication "
+                        f"(eval_set_index={eval_idx} and all TARGET values are NaN)"
+                    )
+                else:
+                    other_dfs.append(eval_subset)
+
+            if oot_eval_dfs:
+                oot_df = pd.concat(oot_eval_dfs, ignore_index=False)
+                df_for_dedup = pd.concat(other_dfs, ignore_index=False)
+            else:
+                df_for_dedup = df
+
+        marked_duplicates = df_for_dedup.duplicated(subset=unique_columns, keep=False)
         if marked_duplicates.sum() > 0:
-            dups_indices = df[marked_duplicates].index.to_list()[:100]
-            nrows_after_tgt_dedup = len(df.drop_duplicates(subset=unique_columns, keep=False))
-            num_dup_rows = nrows_after_full_dedup - nrows_after_tgt_dedup
-            share_tgt_dedup = 100 * num_dup_rows / nrows_after_full_dedup
+            dups_indices = df_for_dedup[marked_duplicates].index.to_list()[:100]
+            nrows_after_tgt_dedup = len(df_for_dedup.drop_duplicates(subset=unique_columns, keep=False))
+            num_dup_rows = len(df_for_dedup) - nrows_after_tgt_dedup
+            share_tgt_dedup = 100 * num_dup_rows / len(df_for_dedup)
 
             msg = bundle.get("dataset_diff_target_duplicates").format(share_tgt_dedup, num_dup_rows, dups_indices)
-            df = df.drop_duplicates(subset=unique_columns, keep=False)
-            logger.info(f"Dataset shape after clean invalid target duplicates: {df.shape}")
+            df_for_dedup = df_for_dedup.drop_duplicates(subset=unique_columns, keep=False)
+            logger.info(f"Dataset shape after clean invalid target duplicates: {df_for_dedup.shape}")
+        # Combine back excluded rows
+        if oot_df is not None:
+            df = pd.concat([df_for_dedup, oot_df], ignore_index=False)
+            marked_duplicates = df.duplicated(subset=unique_columns, keep=False)
+            if marked_duplicates.sum() > 0:
+                dups_indices = df[marked_duplicates].index.to_list()[:100]
+                nrows_after_tgt_dedup = len(df.drop_duplicates(subset=unique_columns, keep=False))
+                num_dup_rows = len(df) - nrows_after_tgt_dedup
+                share_tgt_dedup = 100 * num_dup_rows / len(df)
+                msg = bundle.get("dataset_diff_target_duplicates_oot").format(
+                    share_tgt_dedup, num_dup_rows, dups_indices
+                )
+                df = df.drop_duplicates(subset=unique_columns, keep="first")
+            logger.info(f"Final dataset shape after adding back excluded rows: {df.shape}")
+        else:
+            df = df_for_dedup
 
     return df, msg
 

upgini/utils/feature_info.py

@@ -27,6 +27,7 @@ class FeatureInfo:
     doc_link: str
     data_provider_link: str
     data_source_link: str
+    psi_value: Optional[float] = None
 
     @staticmethod
     def from_metadata(
@@ -47,12 +48,14 @@ class FeatureInfo:
             doc_link=feature_meta.doc_link,
             data_provider_link=feature_meta.data_provider_link,
             data_source_link=feature_meta.data_source_link,
+            psi_value=feature_meta.psi_value,
         )
 
     def to_row(self, bundle: ResourceBundle) -> Dict[str, str]:
         return {
             bundle.get("features_info_name"): self.name,
             bundle.get("features_info_shap"): self.rounded_shap,
+            bundle.get("features_info_psi"): self.psi_value,
             bundle.get("features_info_hitrate"): self.hitrate,
             bundle.get("features_info_value_preview"): self.value_preview,
             bundle.get("features_info_provider"): self.provider,
@@ -64,6 +67,7 @@ class FeatureInfo:
         return {
             bundle.get("features_info_name"): self.internal_name,
             bundle.get("features_info_shap"): self.rounded_shap,
+            bundle.get("features_info_psi"): self.psi_value,
             bundle.get("features_info_hitrate"): self.hitrate,
             bundle.get("features_info_value_preview"): self.value_preview,
             bundle.get("features_info_provider"): self.internal_provider,
@@ -76,6 +80,7 @@ class FeatureInfo:
             bundle.get("features_info_name"): self.internal_name,
             "feature_link": self.doc_link,
             bundle.get("features_info_shap"): self.rounded_shap,
+            bundle.get("features_info_psi"): self.psi_value,
             bundle.get("features_info_hitrate"): self.hitrate,
             bundle.get("features_info_value_preview"): self.value_preview,
             bundle.get("features_info_provider"): self.internal_provider,

upgini/utils/psi.py

@@ -0,0 +1,294 @@
+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,
+    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,
+            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,
+    psi_features_params: StabilityParams = DEFAULT_FEATURES_PARAMS,
+    psi_target_params: StabilityParams = DEFAULT_TARGET_PARAMS,
+) -> Dict[str, float]:
+    empty_res = pd.Series(index=df.columns, data=0)
+
+    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
+
+    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(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:
+            logger.info("Cannot determine target PSI. Skip feature PSI check")
+            return pd.Series(index=df.columns, data=0)
+
+        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(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 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

@@ -5,7 +5,7 @@ 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.target_utils import balance_undersample
 from upgini.utils.ts_utils import get_most_frequent_time_unit, trunc_datetime
@@ -117,6 +117,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 +142,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 +203,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,7 +6,7 @@ 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.metadata import EVAL_SET_INDEX, SYSTEM_RECORD_ID, ModelTaskType
 from upgini.resource_bundle import ResourceBundle, get_custom_bundle, bundle
 from upgini.sampler.random_under_sampler import RandomUnderSampler
 
@@ -132,6 +132,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()