upgini 1.2.89a1__py3-none-any.whl → 1.2.91__py3-none-any.whl

upgini/http.py

@@ -252,6 +252,7 @@ class _RestClient:
 
     # V2
     CHECK_UPLOADED_FILE_URL_FMT_V2 = SERVICE_ROOT_V2 + "search/check-file?fileUploadId={0}"
+    IS_FILE_UPLOADED_URL_FMT_V2 = SERVICE_ROOT_V2 + "search/files/exists?digest={0}"
     INITIAL_SEARCH_URI_FMT_V2 = SERVICE_ROOT_V2 + "search/initial"
     INITIAL_SEARCH_WITHOUT_UPLOAD_URI_FMT_V2 = SERVICE_ROOT_V2 + "search/initial-without-upload?fileUploadId={0}"
     VALIDATION_SEARCH_URI_FMT_V2 = SERVICE_ROOT_V2 + "search/validation?initialSearchTaskId={0}"
@@ -272,6 +273,7 @@ class _RestClient:
     SEARCH_FILE_METADATA_URI_FMT_V2 = SERVICE_ROOT_V2 + "search/{0}/metadata"
     SEARCH_TASK_METADATA_FMT_V3 = SERVICE_ROOT_V2 + "search/metadata-v2/{0}"
     SEARCH_DUMP_INPUT_FMT_V2 = SERVICE_ROOT_V2 + "search/dump-input"
+    SEARCH_DUMP_INPUT_FILE_FMT = SERVICE_ROOT_V2 + "search/dump-input-file"
     TRANSFORM_USAGE_FMT = SERVICE_ROOT_V2 + "user/transform-usage"
 
     UPLOAD_USER_ADS_URI = SERVICE_ROOT + "ads/upload"
@@ -410,32 +412,29 @@ class _RestClient:
         eval_x_path: Optional[str] = None,
         eval_y_path: Optional[str] = None,
     ):
-        api_path = self.SEARCH_DUMP_INPUT_FMT_V2
-        files = {}
-        with open(x_path, "rb") as x_file:
-            files["x"] = ("x.pickle", x_file, "application/octet-stream")
-            if y_path:
-                with open(y_path, "rb") as y_file:
-                    files["y"] = ("y.pickle", y_file, "application/octet-stream")
-                    if eval_x_path and eval_y_path:
-                        with open(eval_x_path, "rb") as eval_x_file, open(eval_y_path, "rb") as eval_y_file:
-                            files["eval_x"] = ("eval_x.pickle", eval_x_file, "application/octet-stream")
-                            files["eval_y"] = ("eval_y.pickle", eval_y_file, "application/octet-stream")
-                            self._with_unauth_retry(
-                                lambda: self._send_post_file_req_v2(
-                                    api_path, files, trace_id=trace_id, need_json_response=False
-                                )
-                            )
-                    else:
-                        self._with_unauth_retry(
-                            lambda: self._send_post_file_req_v2(
-                                api_path, files, trace_id=trace_id, need_json_response=False
-                            )
-                        )
+        api_path = self.SEARCH_DUMP_INPUT_FILE_FMT
+
+        def upload_with_check(path: str, file_name: str):
+            digest_sha256 = self.compute_file_digest(path)
+            if self.is_file_uploaded(trace_id, digest_sha256):
+                # print(f"File {path} was already uploaded with digest {digest_sha256}, skipping")
+                return
             else:
-                self._with_unauth_retry(
-                    lambda: self._send_post_file_req_v2(api_path, files, trace_id=trace_id, need_json_response=False)
-                )
+                with open(path, "rb") as file:
+                    files = {"file": (file_name, file, "application/octet-stream")}
+                    self._with_unauth_retry(
+                        lambda: self._send_post_file_req_v2(
+                            api_path, files, trace_id=trace_id, need_json_response=False
+                        )
+                    )
+
+        upload_with_check(x_path, "x.parquet")
+        if y_path:
+            upload_with_check(y_path, "y.parquet")
+        if eval_x_path:
+            upload_with_check(eval_x_path, "eval_x.parquet")
+        if eval_y_path:
+            upload_with_check(eval_y_path, "eval_y.parquet")
 
     @staticmethod
     def compute_file_digest(filepath: str, algorithm="sha256", chunk_size=4096) -> str:
@@ -514,6 +513,11 @@ class _RestClient:
         )
         return bool(response)
 
+    def is_file_uploaded(self, trace_id: str, digest: str) -> bool:
+        api_path = self.IS_FILE_UPLOADED_URL_FMT_V2.format(digest)
+        response = self._with_unauth_retry(lambda: self._send_get_req(api_path, trace_id))
+        return bool(response)
+
     def initial_search_without_upload_v2(
         self,
         trace_id: str,

upgini/metadata.py

@@ -159,6 +159,9 @@ class ModelTaskType(Enum):
     REGRESSION = "REGRESSION"
     TIMESERIES = "TIMESERIES"
 
+    def is_classification(self) -> bool:
+        return self in [ModelTaskType.BINARY, ModelTaskType.MULTICLASS]
+
 
 class ModelLabelType(Enum):
     GINI = "gini"

upgini/metrics.py

@@ -332,7 +332,7 @@ class EstimatorWrapper:
         self.groups = groups
         self.text_features = text_features
         self.logger = logger or logging.getLogger()
-        self.droped_features = []
+        self.dropped_features = []
         self.converted_to_int = []
         self.converted_to_str = []
         self.converted_to_numeric = []
@@ -381,10 +381,11 @@ class EstimatorWrapper:
         x, y, groups = self._prepare_data(x, y, groups=self.groups)
 
         self.logger.info(f"Before preparing data columns: {x.columns.to_list()}")
-        self.droped_features = []
+        self.dropped_features = []
         self.converted_to_int = []
         self.converted_to_str = []
         self.converted_to_numeric = []
+
         for c in x.columns:
 
             if _get_unique_count(x[c]) < 2:
@@ -392,7 +393,7 @@ class EstimatorWrapper:
                 if c in self.cat_features:
                     self.cat_features.remove(c)
                 x.drop(columns=[c], inplace=True)
-                self.droped_features.append(c)
+                self.dropped_features.append(c)
             elif self.text_features is not None and c in self.text_features:
                 x[c] = x[c].astype(str)
                 self.converted_to_str.append(c)
@@ -427,16 +428,16 @@ class EstimatorWrapper:
                     except (ValueError, TypeError):
                         self.logger.warning(f"Remove feature {c} because it is not numeric and not in cat_features")
                         x.drop(columns=[c], inplace=True)
-                        self.droped_features.append(c)
+                        self.dropped_features.append(c)
 
         return x, y, groups, {}
 
     def _prepare_to_calculate(self, x: pd.DataFrame, y: pd.Series) -> Tuple[pd.DataFrame, np.ndarray, dict]:
         x, y, _ = self._prepare_data(x, y)
 
-        if self.droped_features:
-            self.logger.info(f"Drop features on calculate metrics: {self.droped_features}")
-            x = x.drop(columns=self.droped_features)
+        if self.dropped_features:
+            self.logger.info(f"Drop features on calculate metrics: {self.dropped_features}")
+            x = x.drop(columns=self.dropped_features)
 
         if self.converted_to_int:
             self.logger.info(f"Convert to int features on calculate metrics: {self.converted_to_int}")
@@ -797,7 +798,7 @@ class CatBoostWrapper(EstimatorWrapper):
                     )
                 for f in high_cardinality_features:
                     self.text_features.remove(f)
-                    self.droped_features.append(f)
+                    self.dropped_features.append(f)
                     x = x.drop(columns=f, errors="ignore")
                 return super().cross_val_predict(x, y, baseline_score_column)
             else:
@@ -814,7 +815,7 @@ class CatBoostWrapper(EstimatorWrapper):
                 else:
                     encoded = cat_encoder.transform(x[self.cat_features])
                     cat_features = encoded.columns.to_list()
-                x[self.cat_features] = encoded
+                x.loc[:, self.cat_features] = encoded
             else:
                 cat_features = self.cat_features
 
@@ -897,7 +898,7 @@ class LightGBMWrapper(EstimatorWrapper):
         for c in x.columns:
             if x[c].dtype not in ["category", "int64", "float64", "bool"]:
                 self.logger.warning(f"Feature {c} is not numeric and will be dropped")
-                self.droped_features.append(c)
+                self.dropped_features.append(c)
                 x = x.drop(columns=c, errors="ignore")
         return x, y_numpy, groups, params
 
@@ -988,7 +989,7 @@ class OtherEstimatorWrapper(EstimatorWrapper):
         for c in x.columns:
             if x[c].dtype not in ["category", "int64", "float64", "bool"]:
                 self.logger.warning(f"Feature {c} is not numeric and will be dropped")
-                self.droped_features.append(c)
+                self.dropped_features.append(c)
                 x = x.drop(columns=c, errors="ignore")
         return x, y_numpy, groups, params
 

upgini/resource_bundle/strings.properties

@@ -144,6 +144,7 @@ baseline_score_column_has_na=baseline_score_column contains NaN. Clear it and an
 missing_features_for_transform=Missing some features for transform that were presented on fit: {}
 missing_target_for_transform=Search contains features on target. Please add y to the call and try again
 missing_id_column=Id column {} not found in X: {}
+unknown_id_column_value_in_eval_set=Unknown values in id columns: {}
     # target validation
 empty_target=Target is empty in all rows
 # non_numeric_target=Binary target should be numerical type
@@ -195,6 +196,7 @@ timeseries_invalid_test_size_type=test_size={} should be a float in the (0, 1) r
 timeseries_splits_more_than_samples=Number of splits={} can't be more than number of samples={}
 timeseries_invalid_test_size=Wrong number of samples in a test fold: (test_size * n_samples / n_splits) <= 1
 date_and_id_columns_duplicates=Found {} duplicate rows by date and id_columns. Please remove them and try again
+missing_ids_in_eval_set=Following ids are present in eval set but not in sampled train set: {}. They will be removed from eval set.
     # Upload ads validation
 ads_upload_too_few_rows=At least 1000 records per sample are needed. Increase the sample size for evaluation and resubmit the data
 ads_upload_search_key_not_found=Search key {} wasn't found in dataframe columns

upgini/utils/sample_utils.py

@@ -0,0 +1,416 @@
+from dataclasses import dataclass, field
+import logging
+import numbers
+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.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
+
+
+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 = 200_000
+FIT_SAMPLE_ROWS = 200_000
+FIT_SAMPLE_ROWS_WITH_EVAL_SET = 200_000
+FIT_SAMPLE_THRESHOLD_WITH_EVAL_SET = 200_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
+    target: str
+    ids: Optional[List[str]] = None
+    eval_set_index: Optional[str] = None
+
+
+def sample(
+    df: pd.DataFrame,
+    task_type: Optional[ModelTaskType],
+    cv_type: Optional[CVType],
+    sample_config: SampleConfig,
+    sample_columns: SampleColumns,
+    random_state: int = 42,
+    balance: bool = True,
+    force_downsampling: bool = False,
+    logger: Optional[logging.Logger] = None,
+    **kwargs,
+) -> pd.DataFrame:
+    if force_downsampling:
+        return balance_undersample_forced(
+            df,
+            sample_columns.target,
+            sample_columns.ids,
+            sample_columns.date,
+            task_type,
+            cv_type,
+            random_state,
+            sample_config.force_sample_size,
+            logger=logger,
+            **kwargs,
+        )
+
+    if sample_columns.eval_set_index in df.columns:
+        fit_sample_threshold = sample_config.fit_sample_threshold_with_eval_set
+        fit_sample_rows = sample_config.fit_sample_rows_with_eval_set
+    else:
+        fit_sample_threshold = sample_config.fit_sample_threshold
+        fit_sample_rows = sample_config.fit_sample_rows
+
+    if cv_type is not None and cv_type.is_time_series():
+        return sample_time_series_train_eval(
+            df,
+            sample_columns,
+            sample_config.fit_sample_rows_ts,
+            trim_threshold=fit_sample_threshold,
+            max_rows=fit_sample_rows,
+            random_state=random_state,
+            logger=logger,
+            **kwargs,
+        )
+
+    if task_type is not None and task_type.is_classification() and balance:
+        df = balance_undersample(
+            df=df,
+            target_column=sample_columns.target,
+            task_type=task_type,
+            random_state=random_state,
+            binary_min_sample_threshold=sample_config.binary_min_sample_threshold,
+            multiclass_min_sample_threshold=sample_config.multiclass_min_sample_threshold,
+            binary_bootstrap_loops=sample_config.binary_bootstrap_loops,
+            multiclass_bootstrap_loops=sample_config.multiclass_bootstrap_loops,
+            logger=logger,
+            **kwargs,
+        )
+
+    num_samples = _num_samples(df)
+    if num_samples > fit_sample_threshold:
+        logger.info(
+            f"Etalon has size {num_samples} more than threshold {fit_sample_threshold} "
+            f"and will be downsampled to {fit_sample_rows}"
+        )
+        df = df.sample(n=fit_sample_rows, random_state=random_state)
+        logger.info(f"Shape after threshold resampling: {df.shape}")
+
+    return df
+
+
+def sample_time_series_train_eval(
+    df: pd.DataFrame,
+    sample_columns: SampleColumns,
+    sample_size: int,
+    trim_threshold: int,
+    max_rows: int,
+    random_state: int = 42,
+    logger: Optional[logging.Logger] = None,
+    bundle: Optional[ResourceBundle] = None,
+    **kwargs,
+):
+    if sample_columns.eval_set_index in df.columns:
+        train_df = df[df[sample_columns.eval_set_index] == 0]
+        eval_df = df[df[sample_columns.eval_set_index] > 0]
+    else:
+        train_df = df
+        eval_df = None
+
+    train_df = sample_time_series_trunc(
+        train_df, sample_columns.ids, sample_columns.date, sample_size, random_state, logger=logger, **kwargs
+    )
+    if sample_columns.ids and eval_df is not None:
+        missing_ids = (
+            eval_df[~eval_df[sample_columns.ids].isin(np.unique(train_df[sample_columns.ids]))][sample_columns.ids]
+            .dropna()
+            .drop_duplicates()
+            .values.tolist()
+        )
+        if missing_ids:
+            bundle = bundle or get_custom_bundle()
+            print(bundle.get("missing_ids_in_eval_set").format(missing_ids))
+            eval_df = eval_df.merge(train_df[sample_columns.ids].drop_duplicates())
+
+    if eval_df is not None:
+        if len(eval_df) > trim_threshold - len(train_df):
+            eval_df = sample_time_series_trunc(
+                eval_df,
+                sample_columns.ids,
+                sample_columns.date,
+                max_rows - len(train_df),
+                random_state,
+                logger=logger,
+                **kwargs,
+            )
+        if logger is not None:
+            logger.info(f"Eval set size: {len(eval_df)}")
+        df = pd.concat([train_df, eval_df])
+
+    elif len(train_df) > max_rows:
+        df = sample_time_series_trunc(
+            train_df,
+            sample_columns.ids,
+            sample_columns.date,
+            max_rows,
+            random_state,
+            logger=logger,
+            **kwargs,
+        )
+    else:
+        df = train_df
+
+    if logger is not None:
+        logger.info(f"Train set size: {len(df)}")
+
+    return df
+
+
+def sample_time_series_trunc(
+    df: pd.DataFrame,
+    id_columns: Optional[List[str]],
+    date_column: str,
+    sample_size: int,
+    random_state: int = 42,
+    logger: Optional[logging.Logger] = None,
+    highfreq_trunc_lengths: List[pd.DateOffset] = TS_DEFAULT_HIGH_FREQ_TRUNC_LENGTHS,
+    lowfreq_trunc_lengths: List[pd.DateOffset] = TS_DEFAULT_LOW_FREQ_TRUNC_LENGTHS,
+    time_unit_threshold: pd.Timedelta = TS_DEFAULT_TIME_UNIT_THRESHOLD,
+    **kwargs,
+):
+    if id_columns is None:
+        id_columns = []
+    # Convert date column to datetime
+    dates_df = df[id_columns + [date_column]].copy().reset_index(drop=True)
+    if pd.api.types.is_numeric_dtype(dates_df[date_column]):
+        dates_df[date_column] = pd.to_datetime(dates_df[date_column], unit="ms")
+    else:
+        dates_df[date_column] = pd.to_datetime(dates_df[date_column])
+
+    time_unit = get_most_frequent_time_unit(dates_df, id_columns, date_column)
+    if logger is not None:
+        logger.info(f"Time unit: {time_unit}")
+
+    if time_unit is None:
+        if logger is not None:
+            logger.info("Cannot detect time unit, returning original dataset")
+        return df
+
+    if time_unit < time_unit_threshold:
+        for trunc_length in highfreq_trunc_lengths:
+            sampled_df = trunc_datetime(dates_df, id_columns, date_column, trunc_length, logger=logger)
+            if len(sampled_df) <= sample_size:
+                break
+        if len(sampled_df) > sample_size:
+            sampled_df = sample_time_series(
+                sampled_df, id_columns, date_column, sample_size, random_state, logger=logger, **kwargs
+            )
+    else:
+        for trunc_length in lowfreq_trunc_lengths:
+            sampled_df = trunc_datetime(dates_df, id_columns, date_column, trunc_length, logger=logger)
+            if len(sampled_df) <= sample_size:
+                break
+        if len(sampled_df) > sample_size:
+            sampled_df = sample_time_series(
+                sampled_df, id_columns, date_column, sample_size, random_state, logger=logger, **kwargs
+            )
+
+    return df.iloc[sampled_df.index]
+
+
+def sample_time_series(
+    df: pd.DataFrame,
+    id_columns: List[str],
+    date_column: str,
+    sample_size: int,
+    random_state: int = 42,
+    min_different_ids_ratio: float = TS_MIN_DIFFERENT_IDS_RATIO,
+    prefer_recent_dates: bool = True,
+    logger: Optional[logging.Logger] = None,
+    **kwargs,
+):
+    def ensure_tuple(x):
+        return tuple([x]) if not isinstance(x, tuple) else x
+
+    random_state = np.random.RandomState(random_state)
+
+    if not id_columns:
+        id_columns = [date_column]
+    ids_sort = df.groupby(id_columns)[date_column].aggregate(["max", "count"]).T.to_dict()
+    ids_sort = {
+        ensure_tuple(k): (
+            (v["max"], v["count"], random_state.rand()) if prefer_recent_dates else (v["count"], random_state.rand())
+        )
+        for k, v in ids_sort.items()
+    }
+    id_counts = df[id_columns].value_counts()
+    id_counts.index = [ensure_tuple(i) for i in id_counts.index]
+    id_counts = id_counts.sort_index(key=lambda x: [ids_sort[y] for y in x], ascending=False).cumsum()
+    id_counts = id_counts[id_counts <= sample_size]
+    min_different_ids = max(int(len(df[id_columns].drop_duplicates()) * min_different_ids_ratio), 1)
+
+    def id_mask(sample_index: pd.Index) -> pd.Index:
+        if isinstance(sample_index, pd.MultiIndex):
+            return pd.MultiIndex.from_frame(df[id_columns]).isin(sample_index)
+        else:
+            return df[id_columns[0]].isin(sample_index)
+
+    if len(id_counts) < min_different_ids:
+        if logger is not None:
+            logger.info(
+                f"Different ids count {len(id_counts)} for sample size {sample_size}"
+                f" is less than min different ids {min_different_ids}, sampling time window"
+            )
+        date_counts = df.groupby(id_columns)[date_column].nunique().sort_values(ascending=False)
+        ids_to_sample = date_counts.index[:min_different_ids] if len(id_counts) > 0 else date_counts.index
+        mask = id_mask(ids_to_sample)
+        df = df[mask]
+        sample_date_counts = df[date_column].value_counts().sort_index(ascending=False).cumsum()
+        sample_date_counts = sample_date_counts[sample_date_counts <= sample_size]
+        df = df[df[date_column].isin(sample_date_counts.index)]
+    else:
+        if len(id_columns) > 1:
+            id_counts.index = pd.MultiIndex.from_tuples(id_counts.index)
+        else:
+            id_counts.index = [i[0] for i in id_counts.index]
+        mask = id_mask(id_counts.index)
+        df = df[mask]
+
+    return df
+
+
+def balance_undersample_forced(
+    df: pd.DataFrame,
+    sample_columns: SampleColumns,
+    task_type: ModelTaskType,
+    cv_type: Optional[CVType],
+    random_state: int,
+    sample_size: int = 7000,
+    logger: Optional[logging.Logger] = None,
+    bundle: Optional[ResourceBundle] = None,
+    warning_callback: Optional[Callable] = None,
+):
+    if len(df) <= sample_size:
+        return df
+
+    if logger is None:
+        logger = logging.getLogger("muted_logger")
+        logger.setLevel("FATAL")
+    bundle = bundle or get_custom_bundle()
+    if SYSTEM_RECORD_ID not in df.columns:
+        raise Exception("System record id must be presented for undersampling")
+
+    msg = bundle.get("forced_balance_undersample")
+    logger.info(msg)
+    if warning_callback is not None:
+        warning_callback(msg)
+
+    target = df[sample_columns.target].copy()
+
+    vc = target.value_counts()
+    max_class_value = vc.index[0]
+    min_class_value = vc.index[len(vc) - 1]
+    max_class_count = vc[max_class_value]
+    min_class_count = vc[min_class_value]
+
+    resampled_data = df
+    df = df.copy().sort_values(by=SYSTEM_RECORD_ID)
+    if cv_type is not None and cv_type.is_time_series():
+        logger.warning(f"Sampling time series dataset from {len(df)} to {sample_size}")
+        resampled_data = sample_time_series_train_eval(
+            df,
+            sample_columns=sample_columns,
+            sample_size=sample_size,
+            trim_threshold=sample_size,
+            max_rows=sample_size,
+            random_state=random_state,
+            logger=logger,
+        )
+    elif task_type in [ModelTaskType.MULTICLASS, ModelTaskType.REGRESSION]:
+        logger.warning(f"Sampling dataset from {len(df)} to {sample_size}")
+        resampled_data = df.sample(n=sample_size, random_state=random_state)
+    else:
+        msg = bundle.get("imbalanced_target").format(min_class_value, min_class_count)
+        logger.warning(msg)
+
+        # fill up to min_sample_threshold by majority class
+        minority_class = df[df[sample_columns.target] == min_class_value]
+        majority_class = df[df[sample_columns.target] != min_class_value]
+        logger.info(
+            f"Min class count: {min_class_count}. Max class count: {max_class_count}."
+            f" Rebalance sample size: {sample_size}"
+        )
+        if len(minority_class) > (sample_size / 2):
+            sampled_minority_class = minority_class.sample(n=int(sample_size / 2), random_state=random_state)
+        else:
+            sampled_minority_class = minority_class
+
+        if len(majority_class) > (sample_size) / 2:
+            sampled_majority_class = majority_class.sample(n=int(sample_size / 2), random_state=random_state)
+
+        resampled_data = df[
+            (df[SYSTEM_RECORD_ID].isin(sampled_minority_class[SYSTEM_RECORD_ID]))
+            | (df[SYSTEM_RECORD_ID].isin(sampled_majority_class[SYSTEM_RECORD_ID]))
+        ]
+
+    logger.info(f"Shape after forced rebalance resampling: {resampled_data}")
+    return resampled_data
+
+
+def _num_samples(x):
+    """Return number of samples in array-like x."""
+    if x is None:
+        return 0
+    message = "Expected sequence or array-like, got %s" % type(x)
+    if hasattr(x, "fit") and callable(x.fit):
+        # Don't get num_samples from an ensembles length!
+        raise TypeError(message)
+
+    if not hasattr(x, "__len__") and not hasattr(x, "shape"):
+        if hasattr(x, "__array__"):
+            x = np.asarray(x)
+        else:
+            raise TypeError(message)
+
+    if hasattr(x, "shape") and x.shape is not None:
+        if len(x.shape) == 0:
+            raise TypeError("Singleton array %r cannot be considered a valid collection." % x)
+        # Check that shape is returning an integer or default to len
+        # Dask dataframes may not return numeric shape[0] value
+        if isinstance(x.shape[0], numbers.Integral):
+            return x.shape[0]
+
+    try:
+        return len(x)
+    except TypeError as type_error:
+        raise TypeError(message) from type_error