upgini 1.2.68a3818.dev3__py3-none-any.whl → 1.2.68a3832.dev1__py3-none-any.whl

upgini/__about__.py

@@ -1 +1 @@
-__version__ = "1.2.68a3818.dev3"
+__version__ = "1.2.68a3832.dev1"

upgini/autofe/date.py

@@ -8,7 +8,6 @@ from pandas.core.arrays.timedeltas import TimedeltaArray
 from pydantic import BaseModel, __version__ as pydantic_version
 
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator
-from upgini.autofe.utils import pydantic_validator
 
 
 def get_pydantic_version():
@@ -210,14 +209,6 @@ class DateListDiffBounded(DateListDiff, ParametrizedOperator):
 
         return cls(diff_unit=diff_unit, lower_bound=lower_bound, upper_bound=upper_bound, aggregation=aggregation)
 
-    def get_params(self) -> Dict[str, Optional[str]]:
-        res = super().get_params()
-        if self.lower_bound is not None:
-            res["lower_bound"] = str(self.lower_bound)
-        if self.upper_bound is not None:
-            res["upper_bound"] = str(self.upper_bound)
-        return res
-
     def _agg(self, x):
         x = x[
             (x >= (self.lower_bound if self.lower_bound is not None else -np.inf))
@@ -278,17 +269,32 @@ class DatePercentile(DatePercentileBase):
             {
                 "zero_month": self.zero_month,
                 "zero_year": self.zero_year,
-                "zero_bounds": json.dumps(self.zero_bounds),
+                "zero_bounds": self.zero_bounds,
                 "step": self.step,
             }
         )
         return res
 
-    @pydantic_validator("zero_bounds", mode="before")
-    def parse_zero_bounds(cls, value):
-        if isinstance(value, str):
-            return json.loads(value)
-        return value
+    # Check Pydantic version
+    if get_pydantic_version() >= 2:
+        # Use @field_validator for Pydantic 2.x
+        from pydantic import field_validator
+
+        @field_validator("zero_bounds", mode="before")
+        def parse_zero_bounds(cls, value):
+            if isinstance(value, str):
+                return json.loads(value)
+            return value
+
+    else:
+        # Use @validator for Pydantic 1.x
+        from pydantic import validator
+
+        @validator("zero_bounds", pre=True)
+        def parse_zero_bounds(cls, value):
+            if isinstance(value, str):
+                return json.loads(value)
+            return value
 
     def _get_bounds(self, date_col: pd.Series) -> pd.Series:
         months = date_col.dt.month

upgini/autofe/feature.py

@@ -112,11 +112,7 @@ class Feature:
 
     def get_hash(self) -> str:
         return hashlib.sha256(
-            "_".join(
-                [self.op.get_hash_component()]
-                + [ch.op.get_hash_component() for ch in self.children if isinstance(ch, Feature)]
-                + [ch.get_display_name() for ch in self.children]
-            ).encode("utf-8")
+            "_".join([self.op.get_hash_component()] + [ch.get_display_name() for ch in self.children]).encode("utf-8")
         ).hexdigest()[:8]
 
     def set_alias(self, alias: str) -> "Feature":

upgini/autofe/timeseries/base.py

@@ -85,7 +85,7 @@ class TimeSeriesBase(PandasOperator, abc.ABC):
         """
         import re
 
-        offset_regex = f"{base_regex}_offset_(\\d+)([a-zA-Z])$"
+        offset_regex = f"{base_regex}_offset_(\\d+)([a-zA-Z])"
         match = re.match(offset_regex, formula)
 
         if match:

upgini/autofe/timeseries/cross.py

@@ -1,13 +1,16 @@
-import json
 from typing import Dict, List, Optional
 
 import numpy as np
 import pandas as pd
 
+try:
+    from pydantic import field_validator as validator  # V2
+except ImportError:
+    from pydantic import validator  # V1
+
 from upgini.autofe.all_operators import find_op
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator
 from upgini.autofe.timeseries.base import TimeSeriesBase
-from upgini.autofe.utils import pydantic_validator
 
 
 class CrossSeriesInteraction(TimeSeriesBase, ParametrizedOperator):
@@ -17,22 +20,11 @@ class CrossSeriesInteraction(TimeSeriesBase, ParametrizedOperator):
     left_descriptor: List[str] = []
     right_descriptor: List[str] = []
 
-    @pydantic_validator("descriptor_indices")
+    @validator("descriptor_indices")
+    @classmethod
     def validate_descriptor_indices(cls, v):
-        if isinstance(v, str):
-            return json.loads(v)
-        return v
-
-    @pydantic_validator("left_descriptor", "right_descriptor", mode="before")
-    def parse_descriptors(cls, v):
-        if isinstance(v, str):
-            return json.loads(v)
-        return v
-
-    @pydantic_validator("interaction_op", mode="before")
-    def validate_interaction_op(cls, v):
-        if isinstance(v, str):
-            return find_op(v)
+        if not v:
+            raise ValueError("descriptor_indices cannot be empty for CrossSeriesInteraction")
         return v
 
     def __init__(self, **data):
@@ -96,9 +88,9 @@ class CrossSeriesInteraction(TimeSeriesBase, ParametrizedOperator):
         res.update(
             {
                 "interaction_op": self._get_interaction_op_name(),
-                "descriptor_indices": json.dumps(self.descriptor_indices),
-                "left_descriptor": json.dumps(self.left_descriptor),
-                "right_descriptor": json.dumps(self.right_descriptor),
+                "descriptor_indices": self.descriptor_indices,
+                "left_descriptor": self.left_descriptor,
+                "right_descriptor": self.right_descriptor,
             }
         )
         return res

upgini/autofe/timeseries/roll.py

@@ -3,7 +3,6 @@ from typing import Dict, Optional
 
 from upgini.autofe.operator import ParametrizedOperator
 from upgini.autofe.timeseries.base import TimeSeriesBase
-from upgini.autofe.utils import pydantic_validator
 
 # Roll aggregation functions
 roll_aggregations = {
@@ -13,13 +12,19 @@ roll_aggregations = {
     "iqr": lambda x: x.quantile(0.75) - x.quantile(0.25),
 }
 
+try:
+    from pydantic import field_validator as validator  # V2
+except ImportError:
+    from pydantic import validator  # V1
+
 
 class Roll(TimeSeriesBase, ParametrizedOperator):
     aggregation: str
     window_size: int = 1
     window_unit: str = "D"
 
-    @pydantic_validator("window_unit")
+    @validator("window_unit")
+    @classmethod
     def validate_window_unit(cls, v: str) -> str:
         try:
             pd.tseries.frequencies.to_offset(v)

upgini/features_enricher.py

@@ -702,7 +702,6 @@ class FeaturesEnricher(TransformerMixin):
     def transform(
         self,
         X: pd.DataFrame,
-        y: Optional[pd.Series] = None,
         *args,
         exclude_features_sources: Optional[List[str]] = None,
         keep_input: bool = True,
@@ -767,7 +766,6 @@ class FeaturesEnricher(TransformerMixin):
                 result, _, _ = self.__inner_transform(
                     trace_id,
                     X,
-                    y=y,
                     exclude_features_sources=exclude_features_sources,
                     importance_threshold=importance_threshold,
                     max_features=max_features,
@@ -1684,6 +1682,7 @@ class FeaturesEnricher(TransformerMixin):
                 validated_X,
                 validated_y,
                 eval_set,
+                is_demo_dataset,
                 exclude_features_sources,
                 trace_id,
                 progress_bar,
@@ -1874,147 +1873,158 @@ class FeaturesEnricher(TransformerMixin):
         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:
-        has_eval_set = eval_set is not None
-
-        self.logger.info(f"Transform {'with' if has_eval_set else 'without'} eval_set")
-
-        # Prepare
-        df = self.__combine_train_and_eval_sets(validated_X, validated_y, eval_set)
-        df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
-        df = self.__downsample_for_metrics(df)
-
-        # Transform
+        eval_set_sampled_dict = {}
+        if eval_set is not None:
+            self.logger.info("Transform with eval_set")
+            # concatenate X and eval_set with eval_set_index
+            df = validated_X.copy()
+            df[TARGET] = validated_y
+            df[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 = pd.concat([df, eval_df_with_index])
+
+            df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
+
+            # downsample if need to eval_set threshold
+            num_samples = _num_samples(df)
+            force_downsampling = (
+                not self.disable_force_downsampling
+                and self.columns_for_online_api is not None
+                and num_samples > Dataset.FORCE_SAMPLE_SIZE
+            )
+            # TODO: check that system_record_id was added before this step
+            if force_downsampling:
+                self.logger.info(f"Force downsampling from {num_samples} to {Dataset.FORCE_SAMPLE_SIZE}")
+                df = balance_undersample_forced(
+                    df=df,
+                    target_column=TARGET,
+                    id_columns=self.id_columns,
+                    date_column=self._get_date_column(self.search_keys),
+                    task_type=self.model_task_type,
+                    cv_type=self.cv,
+                    random_state=self.random_state,
+                    sample_size=Dataset.FORCE_SAMPLE_SIZE,
+                    logger=self.logger,
+                    bundle=self.bundle,
+                    warning_callback=self.__log_warning,
+                )
+            elif 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 = df.sample(n=Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS, random_state=self.random_state)
 
-        enriched_df, _, _ = self.__inner_transform(
-            trace_id,
-            X=df.drop(columns=[TARGET]),
-            y=df[TARGET],
-            exclude_features_sources=exclude_features_sources,
-            silent_mode=True,
-            metrics_calculation=True,
-            progress_bar=progress_bar,
-            progress_callback=progress_callback,
-            add_fit_system_record_id=True,
-        )
-        if enriched_df is None:
-            return None
+            eval_set_sampled_dict = {}
 
-        x_columns = [
-            c
-            for c in (validated_X.columns.tolist() + self.fit_generated_features + [SYSTEM_RECORD_ID])
-            if c in enriched_df.columns
-        ]
+            tmp_target_name = "__target"
+            df = df.rename(columns={TARGET: tmp_target_name})
 
-        X_sampled, y_sampled, enriched_X = self.__extract_train_data(enriched_df, x_columns)
-        eval_set_sampled_dict = self.__extract_eval_data(
-            enriched_df, x_columns, enriched_X.columns.tolist(), len(eval_set) if has_eval_set else 0
-        )
+            enriched_df, columns_renaming, generated_features = self.__inner_transform(
+                trace_id,
+                df,
+                exclude_features_sources=exclude_features_sources,
+                silent_mode=True,
+                metrics_calculation=True,
+                progress_bar=progress_bar,
+                progress_callback=progress_callback,
+                add_fit_system_record_id=True,
+                target_name=tmp_target_name,
+            )
+            if enriched_df is None:
+                return None
 
-        # Cache and return results
-        return self.__cache_and_return_results(
-            validated_X, validated_y, eval_set, X_sampled, y_sampled, enriched_X, eval_set_sampled_dict
-        )
+            enriched_df = enriched_df.rename(columns={tmp_target_name: TARGET})
 
-    def __combine_train_and_eval_sets(
-        self, validated_X: pd.DataFrame, validated_y: pd.Series, eval_set: Optional[List[tuple]]
-    ) -> pd.DataFrame:
-        df = validated_X.copy()
-        df[TARGET] = validated_y
-        if eval_set is None:
-            return df
+            x_columns = [
+                c
+                for c in (validated_X.columns.tolist() + generated_features + [SYSTEM_RECORD_ID])
+                if c in enriched_df.columns
+            ]
 
-        df[EVAL_SET_INDEX] = 0
+            enriched_Xy = enriched_df.query(f"{EVAL_SET_INDEX} == 0")
+            X_sampled = enriched_Xy[x_columns].copy()
+            y_sampled = enriched_Xy[TARGET].copy()
+            enriched_X = enriched_Xy.drop(columns=[TARGET, EVAL_SET_INDEX])
+            enriched_X_columns = enriched_X.columns.tolist()
 
-        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 = pd.concat([df, eval_df_with_index])
+            for idx in range(len(eval_set)):
+                enriched_eval_xy = enriched_df.query(f"{EVAL_SET_INDEX} == {idx + 1}")
+                eval_x_sampled = enriched_eval_xy[x_columns].copy()
+                eval_y_sampled = enriched_eval_xy[TARGET].copy()
+                enriched_eval_x = enriched_eval_xy[enriched_X_columns].copy()
+                eval_set_sampled_dict[idx] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
+        else:
+            self.logger.info("Transform without eval_set")
+            df = validated_X.copy()
 
-        return df
+            df[TARGET] = validated_y
 
-    def __downsample_for_metrics(self, df: pd.DataFrame) -> pd.DataFrame:
-        num_samples = _num_samples(df)
-        force_downsampling = (
-            not self.disable_force_downsampling
-            and self.columns_for_online_api is not None
-            and num_samples > Dataset.FORCE_SAMPLE_SIZE
-        )
+            df, _ = clean_full_duplicates(df, logger=self.logger, bundle=self.bundle)
 
-        if force_downsampling:
-            self.logger.info(f"Force downsampling from {num_samples} to {Dataset.FORCE_SAMPLE_SIZE}")
-            return balance_undersample_forced(
-                df=df,
-                target_column=TARGET,
-                id_columns=self.id_columns,
-                date_column=self._get_date_column(self.search_keys),
-                task_type=self.model_task_type,
-                cv_type=self.cv,
-                random_state=self.random_state,
-                sample_size=Dataset.FORCE_SAMPLE_SIZE,
-                logger=self.logger,
-                bundle=self.bundle,
-                warning_callback=self.__log_warning,
+            num_samples = _num_samples(df)
+            force_downsampling = (
+                not self.disable_force_downsampling
+                and self.columns_for_online_api is not None
+                and num_samples > Dataset.FORCE_SAMPLE_SIZE
             )
-        elif num_samples > Dataset.FIT_SAMPLE_THRESHOLD:
-            if EVAL_SET_INDEX in df.columns:
-                threshold = Dataset.FIT_SAMPLE_WITH_EVAL_SET_THRESHOLD
-                sample_size = Dataset.FIT_SAMPLE_WITH_EVAL_SET_ROWS
-            else:
-                threshold = Dataset.FIT_SAMPLE_THRESHOLD
-                sample_size = Dataset.FIT_SAMPLE_ROWS
 
-            if num_samples > threshold:
-                self.logger.info(f"Downsampling from {num_samples} to {sample_size}")
-                return df.sample(n=sample_size, random_state=self.random_state)
+            if force_downsampling:
+                self.logger.info(f"Force downsampling from {num_samples} to {Dataset.FORCE_SAMPLE_SIZE}")
+                df = balance_undersample_forced(
+                    df=df,
+                    target_column=TARGET,
+                    id_columns=self.id_columns,
+                    date_column=self._get_date_column(self.search_keys),
+                    task_type=self.model_task_type,
+                    cv_type=self.cv,
+                    random_state=self.random_state,
+                    sample_size=Dataset.FORCE_SAMPLE_SIZE,
+                    logger=self.logger,
+                    bundle=self.bundle,
+                    warning_callback=self.__log_warning,
+                )
+            elif 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)
 
-        return df
+            tmp_target_name = "__target"
+            df = df.rename(columns={TARGET: tmp_target_name})
 
-    def __extract_train_data(
-        self, enriched_df: pd.DataFrame, x_columns: List[str]
-    ) -> Tuple[pd.DataFrame, pd.Series, pd.DataFrame]:
-        if EVAL_SET_INDEX in enriched_df.columns:
-            enriched_Xy = enriched_df.query(f"{EVAL_SET_INDEX} == 0")
-        else:
-            enriched_Xy = enriched_df
-        X_sampled = enriched_Xy[x_columns].copy()
-        y_sampled = enriched_Xy[TARGET].copy()
-        enriched_X = enriched_Xy.drop(columns=[TARGET, EVAL_SET_INDEX], errors="ignore")
-        return X_sampled, y_sampled, enriched_X
+            enriched_Xy, columns_renaming, generated_features = self.__inner_transform(
+                trace_id,
+                df,
+                exclude_features_sources=exclude_features_sources,
+                silent_mode=True,
+                metrics_calculation=True,
+                progress_bar=progress_bar,
+                progress_callback=progress_callback,
+                add_fit_system_record_id=True,
+                target_name=tmp_target_name,
+            )
+            if enriched_Xy is None:
+                return None
 
-    def __extract_eval_data(
-        self, enriched_df: pd.DataFrame, x_columns: List[str], enriched_X_columns: List[str], eval_set_len: int
-    ) -> Dict[int, Tuple]:
-        eval_set_sampled_dict = {}
+            enriched_Xy = enriched_Xy.rename(columns={tmp_target_name: TARGET})
 
-        for idx in range(eval_set_len):
-            enriched_eval_xy = enriched_df.query(f"{EVAL_SET_INDEX} == {idx + 1}")
-            eval_x_sampled = enriched_eval_xy[x_columns].copy()
-            eval_y_sampled = enriched_eval_xy[TARGET].copy()
-            enriched_eval_x = enriched_eval_xy[enriched_X_columns].copy()
-            eval_set_sampled_dict[idx] = (eval_x_sampled, enriched_eval_x, eval_y_sampled)
+            x_columns = [
+                c
+                for c in (validated_X.columns.tolist() + generated_features + [SYSTEM_RECORD_ID])
+                if c in enriched_Xy.columns
+            ]
 
-        return eval_set_sampled_dict
+            X_sampled = enriched_Xy[x_columns].copy()
+            y_sampled = enriched_Xy[TARGET].copy()
+            enriched_X = enriched_Xy.drop(columns=TARGET)
 
-    def __cache_and_return_results(
-        self,
-        validated_X: pd.DataFrame,
-        validated_y: pd.Series,
-        eval_set: Optional[List[tuple]],
-        X_sampled: pd.DataFrame,
-        y_sampled: pd.Series,
-        enriched_X: pd.DataFrame,
-        eval_set_sampled_dict: Dict[int, Tuple],
-    ) -> _SampledDataForMetrics:
         datasets_hash = hash_input(validated_X, validated_y, eval_set)
-        columns_renaming = getattr(self, "fit_columns_renaming", {})
-
         self.__cached_sampled_datasets[datasets_hash] = (
             X_sampled,
             y_sampled,
@@ -2151,7 +2161,6 @@ if response.status_code == 200:
         trace_id: str,
         X: pd.DataFrame,
         *,
-        y: Optional[pd.Series] = None,
         exclude_features_sources: Optional[List[str]] = None,
         importance_threshold: Optional[float] = None,
         max_features: Optional[int] = None,
@@ -2170,14 +2179,8 @@ if response.status_code == 200:
             self.logger.info("Start transform")
 
             validated_X = self._validate_X(X, is_transform=True)
-            if y is not None:
-                validated_y = self._validate_y(validated_X, y)
-                df = self.__combine_train_and_eval_sets(validated_X, validated_y, eval_set=None)
-            else:
-                validated_y = None
-                df = validated_X
 
-            self.__log_debug_information(validated_X, validated_y, exclude_features_sources=exclude_features_sources)
+            self.__log_debug_information(validated_X, exclude_features_sources=exclude_features_sources)
 
             self.__validate_search_keys(self.search_keys, self.search_id)
 
@@ -2220,27 +2223,29 @@ if response.status_code == 200:
                         self.logger.info(msg)
                         print(msg)
 
-            is_demo_dataset = hash_input(df) in DEMO_DATASET_HASHES
+            is_demo_dataset = hash_input(validated_X) in DEMO_DATASET_HASHES
 
             columns_to_drop = [
-                c for c in df.columns if c in self.feature_names_ and c in self.dropped_client_feature_names_
+                c for c in validated_X.columns if c in self.feature_names_ and c in self.dropped_client_feature_names_
             ]
             if len(columns_to_drop) > 0:
                 msg = self.bundle.get("x_contains_enriching_columns").format(columns_to_drop)
                 self.logger.warning(msg)
                 print(msg)
-                df = df.drop(columns=columns_to_drop)
+                validated_X = validated_X.drop(columns=columns_to_drop)
 
             search_keys = self.search_keys.copy()
             if self.id_columns is not None and self.cv is not None and self.cv.is_time_series():
-                search_keys.update(
+                self.search_keys.update(
                     {col: SearchKey.CUSTOM_KEY for col in self.id_columns if col not in self.search_keys}
                 )
 
             search_keys = self.__prepare_search_keys(
-                df, search_keys, is_demo_dataset, is_transform=True, silent_mode=silent_mode
+                validated_X, search_keys, is_demo_dataset, is_transform=True, silent_mode=silent_mode
             )
 
+            df = validated_X.copy()
+
             df = self.__handle_index_search_keys(df, search_keys)
 
             if DEFAULT_INDEX in df.columns:
@@ -2279,11 +2284,8 @@ if response.status_code == 200:
             features_for_transform = self._search_task.get_features_for_transform() or []
             if len(features_for_transform) > 0:
                 missing_features_for_transform = [
-                    columns_renaming.get(f) or f for f in features_for_transform if f not in df.columns
+                    columns_renaming.get(f) for f in features_for_transform if f not in df.columns
                 ]
-                if TARGET in missing_features_for_transform:
-                    raise ValidationError(self.bundle.get("missing_target_for_transform"))
-
                 if len(missing_features_for_transform) > 0:
                     raise ValidationError(
                         self.bundle.get("missing_features_for_transform").format(missing_features_for_transform)
@@ -2339,10 +2341,11 @@ if response.status_code == 200:
                 converter = PostalCodeSearchKeyConverter(postal_code)
                 df = converter.convert(df)
 
-            meaning_types = {}
-            meaning_types.update({col: FileColumnMeaningType.FEATURE for col in features_for_transform})
-            meaning_types.update({col: key.value for col, key in search_keys.items()})
+            # generated_features = [f for f in generated_features if f in self.fit_generated_features]
 
+            meaning_types = {col: key.value for col, key in search_keys.items()}
+            for col in features_for_transform:
+                meaning_types[col] = FileColumnMeaningType.FEATURE
             features_not_to_pass = [
                 c
                 for c in df.columns
@@ -2351,12 +2354,13 @@ if response.status_code == 200:
                 and c not in [ENTITY_SYSTEM_RECORD_ID, SEARCH_KEY_UNNEST]
             ]
 
-            if add_fit_system_record_id:
+            if add_fit_system_record_id and target_name is not None:
+                reversed_columns_renaming = {v: k for k, v in columns_renaming.items()}
                 df = self.__add_fit_system_record_id(
                     df,
                     search_keys,
                     SYSTEM_RECORD_ID,
-                    TARGET,
+                    reversed_columns_renaming.get(target_name, target_name),
                     columns_renaming,
                     silent=True,
                 )
@@ -4066,7 +4070,10 @@ if response.status_code == 200:
             )
 
         if all(k == SearchKey.CUSTOM_KEY for k in valid_search_keys.values()):
-            msg = self.bundle.get("unregistered_only_personal_keys")
+            if self.__is_registered:
+                msg = self.bundle.get("only_custom_keys")
+            else:
+                msg = self.bundle.get("unregistered_only_personal_keys")
             self.logger.warning(msg + f" Provided search keys: {search_keys}")
             raise ValidationError(msg)
 

upgini/metrics.py

@@ -11,7 +11,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import catboost
 import numpy as np
 import pandas as pd
-from catboost import CatBoost, CatBoostClassifier, CatBoostRegressor, Pool
+from lightgbm import LGBMClassifier, LGBMRegressor
 from numpy import log1p
 from pandas.api.types import is_numeric_dtype
 from sklearn.metrics import check_scoring, get_scorer, make_scorer, roc_auc_score
@@ -88,11 +88,18 @@ CATBOOST_MULTICLASS_PARAMS = {
 
 LIGHTGBM_PARAMS = {
     "random_state": DEFAULT_RANDOM_STATE,
-    "num_leaves": 16,
+    # "num_leaves": 16,
+    # "n_estimators": 150,
+    # "min_child_weight": 1,
     "max_depth": 4,
-    "n_estimators": 150,
+    "max_cat_threshold": 80,
+    "min_data_per_group": 25,
+    "num_boost_round": 150,
+    "cat_l2": 10,
+    "cat_smooth": 12,
     "learning_rate": 0.05,
-    "min_child_weight": 1,
+    "feature_fraction": 1.0,
+    "min_sum_hessian_in_leaf": 0.01,
 }
 
 N_FOLDS = 5
@@ -211,6 +218,14 @@ SUPPORTED_CATBOOST_METRICS = {
 }
 
 
+def is_catboost_estimator(estimator):
+    try:
+        from catboost import CatBoostClassifier, CatBoostRegressor
+        return isinstance(estimator, (CatBoostClassifier, CatBoostRegressor))
+    except ImportError:
+        return False
+
+
 @dataclass
 class _CrossValResults:
     metric: Optional[float]
@@ -351,7 +366,7 @@ class EstimatorWrapper:
                 if shaps is not None:
                     for feature, shap_value in shaps.items():
                         # shap_values_all_folds[feature] = shap_values_all_folds.get(feature, []) + shap_value.tolist()
-                        shap_values_all_folds[feature].extend(shap_value.tolist())
+                        shap_values_all_folds[feature].append(shap_value)
 
         if shap_values_all_folds:
             average_shap_values = {
@@ -431,17 +446,26 @@ class EstimatorWrapper:
             # if metric_name.upper() in SUPPORTED_CATBOOST_METRICS:
             #     params["eval_metric"] = SUPPORTED_CATBOOST_METRICS[metric_name.upper()]
             if target_type == ModelTaskType.MULTICLASS:
-                params = _get_add_params(params, CATBOOST_MULTICLASS_PARAMS)
+                # params = _get_add_params(params, CATBOOST_MULTICLASS_PARAMS)
+                # params = _get_add_params(params, add_params)
+                # estimator = CatBoostWrapper(CatBoostClassifier(**params), **kwargs)
+                params = _get_add_params(params, LIGHTGBM_PARAMS)
                 params = _get_add_params(params, add_params)
-                estimator = CatBoostWrapper(CatBoostClassifier(**params), **kwargs)
+                estimator = LightGBMWrapper(LGBMClassifier(**params), **kwargs)
             elif target_type == ModelTaskType.BINARY:
-                params = _get_add_params(params, CATBOOST_BINARY_PARAMS)
+                # params = _get_add_params(params, CATBOOST_BINARY_PARAMS)
+                # params = _get_add_params(params, add_params)
+                # estimator = CatBoostWrapper(CatBoostClassifier(**params), **kwargs)
+                params = _get_add_params(params, LIGHTGBM_PARAMS)
                 params = _get_add_params(params, add_params)
-                estimator = CatBoostWrapper(CatBoostClassifier(**params), **kwargs)
+                estimator = LightGBMWrapper(LGBMClassifier(**params), **kwargs)
             elif target_type == ModelTaskType.REGRESSION:
-                params = _get_add_params(params, CATBOOST_REGRESSION_PARAMS)
+                # params = _get_add_params(params, CATBOOST_REGRESSION_PARAMS)
+                # params = _get_add_params(params, add_params)
+                # estimator = CatBoostWrapper(CatBoostRegressor(**params), **kwargs)
+                params = _get_add_params(params, LIGHTGBM_PARAMS)
                 params = _get_add_params(params, add_params)
-                estimator = CatBoostWrapper(CatBoostRegressor(**params), **kwargs)
+                estimator = LightGBMWrapper(LGBMRegressor(**params), **kwargs)
             else:
                 raise Exception(bundle.get("metrics_unsupported_target_type").format(target_type))
         else:
@@ -450,7 +474,7 @@ class EstimatorWrapper:
             else:
                 estimator_copy = deepcopy(estimator)
             kwargs["estimator"] = estimator_copy
-            if isinstance(estimator, (CatBoostClassifier, CatBoostRegressor)):
+            if is_catboost_estimator(estimator):
                 if cat_features is not None:
                     for cat_feature in cat_features:
                         if cat_feature not in x.columns:
@@ -458,23 +482,13 @@ class EstimatorWrapper:
                                 f"Client cat_feature `{cat_feature}` not found in x columns: {x.columns.to_list()}"
                             )
                     estimator_copy.set_params(
-                        # cat_features=[x.columns.get_loc(cat_feature) for cat_feature in cat_features]
                         cat_features=cat_features
                     )
                 estimator = CatBoostWrapper(**kwargs)
             else:
-                try:
-                    from lightgbm import LGBMClassifier, LGBMRegressor
-
-                    if isinstance(estimator, (LGBMClassifier, LGBMRegressor)):
-                        estimator = LightGBMWrapper(**kwargs)
-                    else:
-                        logger.warning(
-                            f"Unexpected estimator is used for metrics: {estimator}. "
-                            "Default strategy for category features will be used"
-                        )
-                        estimator = OtherEstimatorWrapper(**kwargs)
-                except ModuleNotFoundError:
+                if isinstance(estimator, (LGBMClassifier, LGBMRegressor)):
+                    estimator = LightGBMWrapper(**kwargs)
+                else:
                     logger.warning(
                         f"Unexpected estimator is used for metrics: {estimator}. "
                         "Default strategy for category features will be used"
@@ -487,7 +501,7 @@ class EstimatorWrapper:
 class CatBoostWrapper(EstimatorWrapper):
     def __init__(
         self,
-        estimator: Union[CatBoostClassifier, CatBoostRegressor],
+        estimator,
         scorer: Callable,
         metric_name: str,
         multiplier: int,
@@ -517,6 +531,7 @@ class CatBoostWrapper(EstimatorWrapper):
         x, y, groups, params = super()._prepare_to_fit(x, y)
 
         # Find embeddings
+        from catboost import CatBoostClassifier
         if hasattr(CatBoostClassifier, "get_embedding_feature_indices"):
             emb_pattern = r"(.+)_emb\d+"
             self.emb_features = [c for c in x.columns if re.match(emb_pattern, c) and is_numeric_dtype(x[c])]
@@ -637,8 +652,9 @@ class CatBoostWrapper(EstimatorWrapper):
             else:
                 raise e
 
-    def calculate_shap(self, x: pd.DataFrame, y: pd.Series, estimator: CatBoost) -> Optional[Dict[str, float]]:
+    def calculate_shap(self, x: pd.DataFrame, y: pd.Series, estimator) -> Optional[Dict[str, float]]:
         try:
+            from catboost import Pool
             # Create Pool for fold data, if need (for example, when categorical features are present)
             fold_pool = Pool(
                 x,
@@ -715,6 +731,34 @@ class LightGBMWrapper(EstimatorWrapper):
             y = correct_string_target(y)
         return x, y, params
 
+    def calculate_shap(self, x: pd.DataFrame, y: pd.Series, estimator) -> Optional[Dict[str, float]]:
+        try:
+            import shap
+            import lightgbm as lgb
+
+            if not isinstance(estimator, (lgb.LGBMRegressor, lgb.LGBMClassifier)):
+                return None
+
+            explainer = shap.TreeExplainer(estimator)
+
+            shap_values = explainer.shap_values(x)
+
+            # For classification, shap_values is returned as a list for each class
+            # Take values for the positive class
+            if isinstance(shap_values, list):
+                shap_values = shap_values[1]
+
+            # Calculate mean absolute SHAP value for each feature
+            feature_importance = {}
+            for i, col in enumerate(x.columns):
+                feature_importance[col] = np.mean(np.abs(shap_values[:, i]))
+
+            return feature_importance
+
+        except Exception as e:
+            self.logger.warning(f"Failed to calculate SHAP values: {str(e)}")
+            return None
+
 
 class OtherEstimatorWrapper(EstimatorWrapper):
     def __init__(

upgini/resource_bundle/strings.properties

@@ -80,6 +80,7 @@ email_and_hem_simultanious=EMAIL and HEM search keys cannot be used simultaneous
 postal_code_without_country=COUNTRY search key required if POSTAL_CODE is present
 multiple_search_key=Search key {} passed multiple times
 unregistered_only_personal_keys=Only personal search keys used. Api_key from profile.upgini.com required for EMAIL/HEM, PHONE NUMBER or IPv4/IPv6 search keys\nSee docs https://github.com/upgini/upgini#-open-up-all-capabilities-of-upgini
+only_custom_keys=Only CUSTOM_KEY search keys were provided. At least one of DATE, COUNTRY, POSTAL_CODE, PHONE, EMAIL, HEM, IP should be provided
 search_key_not_found=Column `{}` from search_keys was not found in X dataframe: {}
 numeric_search_key_not_found=Index {} in search_keys is out of bounds for {} columns of X dataframe
 unsupported_search_key_type=Unsupported type of key in search_keys: {}
@@ -136,7 +137,6 @@ x_and_eval_x_diff_types=X and eval_set X has different types: {} and {}
 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
 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
     # target validation
 empty_target=Target is empty in all rows

upgini/search_task.py

@@ -168,13 +168,7 @@ class SearchTask:
         for meta in self.provider_metadata_v2:
             if meta.features_used_for_embeddings is not None:
                 features_for_transform.update(meta.features_used_for_embeddings)
-            if meta.generated_features:
-                features_for_transform.update(
-                    c.original_name
-                    for f in meta.generated_features
-                    for c in f.base_columns
-                    if c.ads_definition_id is None
-                )
+
         return list(features_for_transform)
 
     def get_shuffle_kfold(self) -> Optional[bool]:

upgini/utils/mstats.py

@@ -118,7 +118,7 @@ def spearmanr(
     # - dof: degrees of freedom
     # - t_stat: t-statistic
     # - alternative: 'two-sided', 'greater', 'less'
-    def compute_t_pvalue(t_stat, dof, alternative="two-sided"):
+    def compute_t_pvalue(t_stat, dof, alternative='two-sided'):
         from scipy.stats import t
 
         if alternative == "two-sided":

{upgini-1.2.68a3818.dev3.dist-info → upgini-1.2.68a3832.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.68a3818.dev3
+Version: 1.2.68a3832.dev1
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/
@@ -23,11 +23,11 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Information Analysis
 Requires-Python: <3.12,>=3.8
-Requires-Dist: catboost>=1.0.3
 Requires-Dist: fastparquet>=0.8.1
 Requires-Dist: ipywidgets>=8.1.0
 Requires-Dist: jarowinkler>=2.0.0
 Requires-Dist: levenshtein>=0.25.1
+Requires-Dist: lightgbm>=4.6.0
 Requires-Dist: numpy<=1.26.4,>=1.19.0
 Requires-Dist: pandas<3.0.0,>=1.1.0
 Requires-Dist: psutil>=6.0.0
@@ -39,6 +39,7 @@ Requires-Dist: python-json-logger>=3.3.0
 Requires-Dist: requests>=2.8.0
 Requires-Dist: scikit-learn>=1.3.0
 Requires-Dist: scipy>=1.10.0
+Requires-Dist: shap>=0.44.0
 Requires-Dist: xhtml2pdf<0.3.0,>=0.2.11
 Description-Content-Type: text/markdown
 

{upgini-1.2.68a3818.dev3.dist-info → upgini-1.2.68a3832.dev1.dist-info}/RECORD

@@ -1,14 +1,14 @@
-upgini/__about__.py,sha256=8EJXVXOso85-DEaqFexyMttcmIxDfu1mNawyiRqzDn0,33
+upgini/__about__.py,sha256=Rima-NAKm_MGtHsSB_1p736zBApReFwNHjULA-cfSyg,33
 upgini/__init__.py,sha256=LXSfTNU0HnlOkE69VCxkgIKDhWP-JFo_eBQ71OxTr5Y,261
 upgini/ads.py,sha256=nvuRxRx5MHDMgPr9SiU-fsqRdFaBv8p4_v1oqiysKpc,2714
 upgini/dataset.py,sha256=OGjpeFHbj3lWiZTOHTpWEoMMDmFY1FlNC44FKktoZvU,34956
 upgini/errors.py,sha256=2b_Wbo0OYhLUbrZqdLIx5jBnAsiD1Mcenh-VjR4HCTw,950
-upgini/features_enricher.py,sha256=KBTdADF7_Wj3uDROYdevukOk6R8LVQw47gJkH4M1_iQ,204435
+upgini/features_enricher.py,sha256=GXXx14jwf3F26_KrfJ6O40Vcu1hRx5iBjUB_jxy3Xvg,205476
 upgini/http.py,sha256=ud0Cp7h0jNeHuuZGpU_1dAAEiabGoJjGxc1X5oeBQr4,43496
 upgini/lazy_import.py,sha256=74gQ8JuA48BGRLxAo7lNHNKY2D2emMxrUxKGdxVGhuY,1012
 upgini/metadata.py,sha256=Jh6YTaS00m_nbaOY_owvlSyn9zgkErkqu8iTr9ZjKI8,12279
-upgini/metrics.py,sha256=t7uOOnlDYvP6E3DLjPMQcFBjyhJfUQY8aUlx7N0Mh-s,35477
-upgini/search_task.py,sha256=EuCGp0iCWz2fpuJgN6M47aP_CtIi3Oq9zw78w0mkKiU,17595
+upgini/metrics.py,sha256=OW2a3UWdMEkhRv7XDJvgBsc3iU6RLC5mtqvT1fLURwk,36983
+upgini/search_task.py,sha256=qxUxAD-bed-FpZYmTB_4orW7YJsW_O6a1TcgnZIRFr4,17307
 upgini/spinner.py,sha256=4iMd-eIe_BnkqFEMIliULTbj6rNI2HkN_VJ4qYe0cUc,1118
 upgini/version_validator.py,sha256=DvbaAvuYFoJqYt0fitpsk6Xcv-H1BYDJYHUMxaKSH_Y,1509
 upgini/ads_management/__init__.py,sha256=qzyisOToVRP-tquAJD1PblZhNtMrOB8FiyF9JvfkvgE,50
@@ -16,19 +16,18 @@ upgini/ads_management/ads_manager.py,sha256=igVbN2jz80Umb2BUJixmJVj-zx8unoKpecVo
 upgini/autofe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 upgini/autofe/all_operators.py,sha256=rdjF5eaE4bC6Q4eu_el5Z7ekYt8DjOFermz2bePPbUc,333
 upgini/autofe/binary.py,sha256=MnQuFiERpocjCPQUjOljlsq5FE-04GPfwtNjzvfNMyU,7671
-upgini/autofe/date.py,sha256=C86F7sPiscUGq2a45UtQA9ADWBWg0kt54mePHHzjbLE,10633
-upgini/autofe/feature.py,sha256=y1x3wijhTVBmloayQAHiscqKU9Ll8kLcGm1PdvS357I,14910
+upgini/autofe/date.py,sha256=I07psJerrxOcHao91PdSCk9X6KWu61IBVyFRLjGNgK8,10730
+upgini/autofe/feature.py,sha256=xgu6bVIlUJ5PCUgoXQRNcGkcMOhj-_BdDRmkB_qRFS4,14766
 upgini/autofe/groupby.py,sha256=IYmQV9uoCdRcpkeWZj_kI3ObzoNCNx3ff3h8sTL01tk,3603
 upgini/autofe/operator.py,sha256=EOffJw6vKXpEh5yymqb1RFNJPxGxmnHdFRo9dB5SCFo,4969
 upgini/autofe/unary.py,sha256=yVgPvtfnPSOhrii0YgezddmgWPwyOBCR0JutaIkdTTc,4658
-upgini/autofe/utils.py,sha256=fK1am2_tQj3fL2vDslblye8lmyfWgGIUOX1beYVBz4k,2420
 upgini/autofe/vector.py,sha256=l0KdKg-txlZxDSE4hPPfCtfGQofYbl7oaABPr830sPI,667
 upgini/autofe/timeseries/__init__.py,sha256=PGwwDAMwvkXl3el12tXVEmZUgDUvlmIPlXtROm6bD18,738
-upgini/autofe/timeseries/base.py,sha256=-bzVQ3YgZ5FkRZ64Pm1_fZugaKCixBIDPxOvo6pI370,3756
-upgini/autofe/timeseries/cross.py,sha256=M3aKc_yKSG3Q1xpoTUd51K8kuLYLBJepw9lcvWZCF5Y,5219
+upgini/autofe/timeseries/base.py,sha256=T9Ec8LKJbiwTUGGsd_xhM0U0NUJblqmKchkzUI1sK88,3755
+upgini/autofe/timeseries/cross.py,sha256=Sh5hAXZFWKaFRqf_JGODu9pWO2tmuV5VKyK9eX3i7-I,4931
 upgini/autofe/timeseries/delta.py,sha256=h0YhmI1TlPJnjwFpN_GQxLb6r59DQuucnG5tQAXSgjU,3520
 upgini/autofe/timeseries/lag.py,sha256=LfQtg484vuqM0mgY4Wft1swHX_Srq7OKKgZswCXoiXI,1882
-upgini/autofe/timeseries/roll.py,sha256=zADKXU-eYWQnQ5R3am1yEal8uU6Tm0jLAixwPb_aCHg,2794
+upgini/autofe/timeseries/roll.py,sha256=bNFMDszSYTWvB7EyhHbRY1DJqzSURvHlPAcBebt0y0Y,2878
 upgini/autofe/timeseries/trend.py,sha256=9p2Q5ByAi6cx9RH9teBTe8FyjSzqthznC2Lo5dsJ0ho,2051
 upgini/autofe/timeseries/volatility.py,sha256=9shUmIKjpWTHVYjj80YBsk0XheBJ9uBuLv5NW9Mchnk,7953
 upgini/data_source/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -39,7 +38,7 @@ upgini/normalizer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU
 upgini/normalizer/normalize_utils.py,sha256=Ft2MwSgVoBilXAORAOYAuwPD79GOLfwn4qQE3IUFzzg,7218
 upgini/resource_bundle/__init__.py,sha256=S5F2G47pnJd2LDpmFsjDqEwiKkP8Hm-hcseDbMka6Ko,8345
 upgini/resource_bundle/exceptions.py,sha256=5fRvx0_vWdE1-7HcSgF0tckB4A9AKyf5RiinZkInTsI,621
-upgini/resource_bundle/strings.properties,sha256=LDT-jtYlrD1IXvWjFSf-dtvapje0qSrqI9W3v7y2zVo,27646
+upgini/resource_bundle/strings.properties,sha256=XU5ulr5ZDQfGbFk9QdFDzl3oDMaw0eDYCPoEq3ZvIkw,27687
 upgini/resource_bundle/strings_widget.properties,sha256=gOdqvZWntP2LCza_tyVk1_yRYcG4c04K9sQOAVhF_gw,1577
 upgini/sampler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 upgini/sampler/base.py,sha256=7GpjYqjOp58vYcJLiX__1R5wjUlyQbxvHJ2klFnup_M,6389
@@ -61,7 +60,7 @@ upgini/utils/feature_info.py,sha256=m1tQcT3hTChPAiXzpk0WQcEqElj8KgeCifEJFa7-gss,
 upgini/utils/features_validator.py,sha256=lEfmk4DoxZ4ooOE1HC0ZXtUb_lFKRFHIrnFULZ4_rL8,3746
 upgini/utils/format.py,sha256=Yv5cvvSs2bOLUzzNu96Pu33VMDNbabio92QepUj41jU,243
 upgini/utils/ip_utils.py,sha256=TSQ_qDsLlVnm09X1HacpabEf_HNqSWpxBF4Sdc2xs08,6580
-upgini/utils/mstats.py,sha256=u3gQVUtDRbyrOQK6V1UJ2Rx1QbkSNYGjXa6m3Z_dPVs,6286
+upgini/utils/mstats.py,sha256=dLJQr5Ak5BAoV-pDPpnfvMURZVkZ3_v250QzAsSlqY4,6286
 upgini/utils/phone_utils.py,sha256=IrbztLuOJBiePqqxllfABWfYlfAjYevPhXKipl95wUI,10432
 upgini/utils/postal_code_utils.py,sha256=5M0sUqH2DAr33kARWCTXR-ACyzWbjDq_-0mmEml6ZcU,1716
 upgini/utils/progress_bar.py,sha256=N-Sfdah2Hg8lXP_fV9EfUTXz_PyRt4lo9fAHoUDOoLc,1550
@@ -71,7 +70,7 @@ upgini/utils/target_utils.py,sha256=b1GzO8_gMcwXSZ2v98CY50MJJBzKbWHId_BJGybXfkM,
 upgini/utils/track_info.py,sha256=G5Lu1xxakg2_TQjKZk4b5SvrHsATTXNVV3NbvWtT8k8,5663
 upgini/utils/ts_utils.py,sha256=26vhC0pN7vLXK6R09EEkMK3Lwb9IVPH7LRdqFIQ3kPs,1383
 upgini/utils/warning_counter.py,sha256=-GRY8EUggEBKODPSuXAkHn9KnEQwAORC0mmz_tim-PM,254
-upgini-1.2.68a3818.dev3.dist-info/METADATA,sha256=3Enx8SVCc3I6fz6BmtKXgXiQnrxAUg9ONqYTTTWogTM,49123
-upgini-1.2.68a3818.dev3.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-upgini-1.2.68a3818.dev3.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
-upgini-1.2.68a3818.dev3.dist-info/RECORD,,
+upgini-1.2.68a3832.dev1.dist-info/METADATA,sha256=8eYvZ97d0FtJE5Vj6-AQOm6nNxCHrq2dMggJoF0ft7g,49151
+upgini-1.2.68a3832.dev1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+upgini-1.2.68a3832.dev1.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
+upgini-1.2.68a3832.dev1.dist-info/RECORD,,

{upgini-1.2.68a3818.dev3.dist-info → upgini-1.2.68a3832.dev1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.24.2
+Generator: hatchling 1.25.0
 Root-Is-Purelib: true
 Tag: py3-none-any

upgini/autofe/utils.py

@@ -1,83 +0,0 @@
-"""
-Utility functions for autofe module.
-"""
-
-import functools
-from typing import Callable
-
-
-def get_pydantic_version():
-    """
-    Get the major version of pydantic.
-
-    Returns:
-        int: Major version number (1 or 2)
-    """
-    try:
-        from pydantic import __version__ as pydantic_version
-
-        major_version = int(pydantic_version.split(".")[0])
-        return major_version
-    except (ImportError, ValueError):
-        # Default to version 1 if unable to determine
-        return 1
-
-
-def pydantic_validator(field_name: str, *fields, mode: str = "before", **kwargs):
-    """
-    A decorator that applies the appropriate Pydantic validator based on the installed version.
-
-    This decorator handles the differences between Pydantic v1 and v2 validator syntax,
-    making it easier to write code that works with both versions.
-
-    Args:
-        field_name (str): The name of the field to validate
-        mode (str): The validation mode, either "before" or "after" (for Pydantic v2)
-        **kwargs: Additional arguments to pass to the validator
-
-    Returns:
-        Callable: A decorator that can be applied to validator methods
-
-    Example:
-        ```python
-        class MyModel(BaseModel):
-            items: List[int]
-
-            @pydantic_validator("items")
-            def parse_items(cls, value):
-                if isinstance(value, str):
-                    return [int(x) for x in value.split(",")]
-                return value
-        ```
-    """
-    pydantic_version = get_pydantic_version()
-
-    if pydantic_version >= 2:
-        # Use field_validator for Pydantic 2.x
-        from pydantic import field_validator
-
-        def decorator(func: Callable) -> Callable:
-            @field_validator(field_name, *fields, mode=mode, **kwargs)
-            @functools.wraps(func)
-            def wrapper(cls, value, **kw):
-                return func(cls, value)
-
-            return wrapper
-
-        return decorator
-    else:
-        # Use validator for Pydantic 1.x
-        from pydantic import validator
-
-        # Map mode to Pydantic v1 parameters
-        pre = True if mode == "before" else False
-
-        def decorator(func: Callable) -> Callable:
-            @validator(field_name, *fields, pre=pre, **kwargs)
-            @functools.wraps(func)
-            def wrapper(cls, value, **kw):
-                return func(cls, value)
-
-            return wrapper
-
-        return decorator