upgini 1.1.316a3__py3-none-any.whl → 1.1.316a5__py3-none-any.whl

upgini/__about__.py

@@ -1 +1 @@
-__version__ = "1.1.316a3"
+__version__ = "1.1.316a5"

upgini/autofe/date.py

@@ -4,11 +4,16 @@ from typing import Any, Dict, List, Optional, Union
 import numpy as np
 import pandas as pd
 from pandas.core.arrays.timedeltas import TimedeltaArray
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, __version__ as pydantic_version
 
 from upgini.autofe.operand import PandasOperand
 
 
+def get_pydantic_version():
+    major_version = int(pydantic_version.split('.')[0])
+    return major_version
+
+
 class DateDiffMixin(BaseModel):
     diff_unit: str = "D"
     left_unit: Optional[str] = None
@@ -246,12 +251,25 @@ class DatePercentile(DatePercentileBase):
         )
         return res
 
-    @validator("zero_bounds", pre="true")
-    def validate_bounds(cls, value):
-        if value is None or isinstance(value, list):
+    # 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 value[1:-1].split(", ")
+            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 value[1:-1].split(", ")
             return value
-        elif isinstance(value, str):
-            return value[1:-1].split(", ")
 
     def _get_bounds(self, date_col: pd.Series) -> pd.Series:
         months = date_col.dt.month

upgini/autofe/feature.py

@@ -82,9 +82,9 @@ class Feature:
         self.alias = alias
 
     def set_op_params(self, params: Optional[Dict[str, str]]) -> "Feature":
-        obj_dict = self.op.model_dump().copy()
+        obj_dict = self.op.dict().copy()
         obj_dict.update(params or {})
-        self.op = self.op.__class__.model_validate(obj_dict)
+        self.op = self.op.__class__.parse_obj(obj_dict)
         self.op.set_params(params)
 
         for child in self.children:

upgini/dataset.py

@@ -18,6 +18,7 @@ from pandas.api.types import (
 from upgini.errors import ValidationError
 from upgini.http import ProgressStage, SearchProgress, _RestClient
 from upgini.metadata import (
+    ENTITY_SYSTEM_RECORD_ID,
     EVAL_SET_INDEX,
     SYSTEM_RECORD_ID,
     TARGET,
@@ -157,7 +158,11 @@ class Dataset:  # (pd.DataFrame):
             raise ValidationError(self.bundle.get("dataset_too_few_rows").format(self.MIN_ROWS_COUNT))
 
     def __validate_max_row_count(self):
-        if len(self.data) > self.MAX_ROWS:
+        if ENTITY_SYSTEM_RECORD_ID in self.data.columns:
+            rows_count = self.data[ENTITY_SYSTEM_RECORD_ID].nunique()
+        else:
+            rows_count = len(self.data)
+        if rows_count > self.MAX_ROWS:
             raise ValidationError(self.bundle.get("dataset_too_many_rows_registered").format(self.MAX_ROWS))
 
     def __target_value(self) -> pd.Series:
@@ -199,14 +204,14 @@ class Dataset:  # (pd.DataFrame):
         elif self.task_type == ModelTaskType.REGRESSION:
             if not is_float_dtype(target):
                 try:
-                    self.data[target_column] = self.data[target_column].astype("float")
+                    self.data[target_column] = self.data[target_column].astype("float64")
                 except ValueError:
                     self.logger.exception("Failed to cast target to float for regression task type")
                     raise ValidationError(self.bundle.get("dataset_invalid_regression_target").format(target.dtype))
         elif self.task_type == ModelTaskType.TIMESERIES:
             if not is_float_dtype(target):
                 try:
-                    self.data[target_column] = self.data[target_column].astype("float")
+                    self.data[target_column] = self.data[target_column].astype("float64")
                 except ValueError:
                     self.logger.exception("Failed to cast target to float for timeseries task type")
                     raise ValidationError(self.bundle.get("dataset_invalid_timeseries_target").format(target.dtype))

upgini/features_enricher.py

@@ -2042,7 +2042,7 @@ class FeaturesEnricher(TransformerMixin):
 
             df[ENTITY_SYSTEM_RECORD_ID] = pd.util.hash_pandas_object(
                 df[columns_for_system_record_id], index=False
-            ).astype("Float64")
+            ).astype("float64")
 
             # Explode multiple search keys
             df, unnest_search_keys = self._explode_multiple_search_keys(df, search_keys, columns_renaming)
@@ -2108,7 +2108,7 @@ class FeaturesEnricher(TransformerMixin):
             # search keys might be changed after explode
             columns_for_system_record_id = sorted(list(search_keys.keys()) + features_for_transform)
             df[SYSTEM_RECORD_ID] = pd.util.hash_pandas_object(df[columns_for_system_record_id], index=False).astype(
-                "Float64"
+                "float64"
             )
             meaning_types[SYSTEM_RECORD_ID] = FileColumnMeaningType.SYSTEM_RECORD_ID
             meaning_types[ENTITY_SYSTEM_RECORD_ID] = FileColumnMeaningType.ENTITY_SYSTEM_RECORD_ID

upgini/http.py

@@ -440,18 +440,18 @@ class _RestClient:
                 content = file.read()
                 md5_hash.update(content)
                 digest = md5_hash.hexdigest()
-                metadata_with_md5 = metadata.model_copy(update={"checksumMD5": digest})
+                metadata_with_md5 = metadata.copy(update={"checksumMD5": digest})
 
             digest_sha256 = hashlib.sha256(
                 pd.util.hash_pandas_object(pd.read_parquet(file_path, engine="fastparquet")).values
             ).hexdigest()
-            metadata_with_md5 = metadata_with_md5.model_copy(update={"digest": digest_sha256})
+            metadata_with_md5 = metadata_with_md5.copy(update={"digest": digest_sha256})
 
             with open(file_path, "rb") as file:
                 files = {
                     "metadata": (
                         "metadata.json",
-                        metadata_with_md5.model_dump_json(exclude_none=True).encode(),
+                        metadata_with_md5.json(exclude_none=True).encode(),
                         "application/json",
                     ),
                     "tracking": (
@@ -461,7 +461,7 @@ class _RestClient:
                     ),
                     "metrics": (
                         "metrics.json",
-                        metrics.model_dump_json(exclude_none=True).encode(),
+                        metrics.json(exclude_none=True).encode(),
                         "application/json",
                     ),
                     "file": (metadata_with_md5.name, file, "application/octet-stream"),
@@ -469,7 +469,7 @@ class _RestClient:
                 if search_customization is not None:
                     files["customization"] = (
                         "customization.json",
-                        search_customization.model_dump_json(exclude_none=True).encode(),
+                        search_customization.json(exclude_none=True).encode(),
                         "application/json",
                     )
                 additional_headers = {self.SEARCH_KEYS_HEADER_NAME: ",".join(self.search_keys_meaning_types(metadata))}
@@ -484,7 +484,7 @@ class _RestClient:
     def check_uploaded_file_v2(self, trace_id: str, file_upload_id: str, metadata: FileMetadata) -> bool:
         api_path = self.CHECK_UPLOADED_FILE_URL_FMT_V2.format(file_upload_id)
         response = self._with_unauth_retry(
-            lambda: self._send_post_req(api_path, trace_id, metadata.model_dump_json(exclude_none=True))
+            lambda: self._send_post_req(api_path, trace_id, metadata.json(exclude_none=True))
         )
         return bool(response)
 
@@ -498,11 +498,11 @@ class _RestClient:
     ) -> SearchTaskResponse:
         api_path = self.INITIAL_SEARCH_WITHOUT_UPLOAD_URI_FMT_V2.format(file_upload_id)
         files = {
-            "metadata": ("metadata.json", metadata.model_dump_json(exclude_none=True).encode(), "application/json"),
-            "metrics": ("metrics.json", metrics.model_dump_json(exclude_none=True).encode(), "application/json"),
+            "metadata": ("metadata.json", metadata.json(exclude_none=True).encode(), "application/json"),
+            "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
         }
         if search_customization is not None:
-            files["customization"] = search_customization.model_dump_json(exclude_none=True).encode()
+            files["customization"] = search_customization.json(exclude_none=True).encode()
         additional_headers = {self.SEARCH_KEYS_HEADER_NAME: ",".join(self.search_keys_meaning_types(metadata))}
         response = self._with_unauth_retry(
             lambda: self._send_post_file_req_v2(
@@ -528,18 +528,18 @@ class _RestClient:
                 content = file.read()
                 md5_hash.update(content)
                 digest = md5_hash.hexdigest()
-                metadata_with_md5 = metadata.model_copy(update={"checksumMD5": digest})
+                metadata_with_md5 = metadata.copy(update={"checksumMD5": digest})
 
             digest_sha256 = hashlib.sha256(
                 pd.util.hash_pandas_object(pd.read_parquet(file_path, engine="fastparquet")).values
             ).hexdigest()
-            metadata_with_md5 = metadata_with_md5.model_copy(update={"digest": digest_sha256})
+            metadata_with_md5 = metadata_with_md5.copy(update={"digest": digest_sha256})
 
             with open(file_path, "rb") as file:
                 files = {
                     "metadata": (
                         "metadata.json",
-                        metadata_with_md5.model_dump_json(exclude_none=True).encode(),
+                        metadata_with_md5.json(exclude_none=True).encode(),
                         "application/json",
                     ),
                     "tracking": (
@@ -549,7 +549,7 @@ class _RestClient:
                     ),
                     "metrics": (
                         "metrics.json",
-                        metrics.model_dump_json(exclude_none=True).encode(),
+                        metrics.json(exclude_none=True).encode(),
                         "application/json",
                     ),
                     "file": (metadata_with_md5.name, file, "application/octet-stream"),
@@ -557,7 +557,7 @@ class _RestClient:
                 if search_customization is not None:
                     files["customization"] = (
                         "customization.json",
-                        search_customization.model_dump_json(exclude_none=True).encode(),
+                        search_customization.json(exclude_none=True).encode(),
                         "application/json",
                     )
 
@@ -581,11 +581,11 @@ class _RestClient:
     ) -> SearchTaskResponse:
         api_path = self.VALIDATION_SEARCH_WITHOUT_UPLOAD_URI_FMT_V2.format(file_upload_id, initial_search_task_id)
         files = {
-            "metadata": ("metadata.json", metadata.model_dump_json(exclude_none=True).encode(), "application/json"),
-            "metrics": ("metrics.json", metrics.model_dump_json(exclude_none=True).encode(), "application/json"),
+            "metadata": ("metadata.json", metadata.json(exclude_none=True).encode(), "application/json"),
+            "metrics": ("metrics.json", metrics.json(exclude_none=True).encode(), "application/json"),
         }
         if search_customization is not None:
-            files["customization"] = search_customization.model_dump_json(exclude_none=True).encode()
+            files["customization"] = search_customization.json(exclude_none=True).encode()
         additional_headers = {self.SEARCH_KEYS_HEADER_NAME: ",".join(self.search_keys_meaning_types(metadata))}
         response = self._with_unauth_retry(
             lambda: self._send_post_file_req_v2(
@@ -649,7 +649,7 @@ class _RestClient:
                     "file": (metadata.name, file, "application/octet-stream"),
                     "metadata": (
                         "metadata.json",
-                        metadata.model_dump_json(exclude_none=True).encode(),
+                        metadata.json(exclude_none=True).encode(),
                         "application/json",
                     ),
                 }
@@ -661,12 +661,12 @@ class _RestClient:
     def get_search_file_metadata(self, search_task_id: str, trace_id: str) -> FileMetadata:
         api_path = self.SEARCH_FILE_METADATA_URI_FMT_V2.format(search_task_id)
         response = self._with_unauth_retry(lambda: self._send_get_req(api_path, trace_id))
-        return FileMetadata.model_validate(response)
+        return FileMetadata.parse_obj(response)
 
     def get_provider_search_metadata_v3(self, provider_search_task_id: str, trace_id: str) -> ProviderTaskMetadataV2:
         api_path = self.SEARCH_TASK_METADATA_FMT_V3.format(provider_search_task_id)
         response = self._with_unauth_retry(lambda: self._send_get_req(api_path, trace_id))
-        return ProviderTaskMetadataV2.model_validate(response)
+        return ProviderTaskMetadataV2.parse_obj(response)
 
     def get_current_transform_usage(self, trace_id) -> TransformUsage:
         track_metrics = get_track_metrics(self.client_ip, self.client_visitorid)

upgini/utils/postal_code_utils.py

@@ -25,7 +25,7 @@ class PostalCodeSearchKeyConverter:
         if is_string_dtype(df[self.postal_code_column]) or is_object_dtype(df[self.postal_code_column]):
             try:
                 df[self.postal_code_column] = (
-                    df[self.postal_code_column].astype("string").astype("Float64").astype("Int64").astype("string")
+                    df[self.postal_code_column].astype("string").astype("float64").astype("Int64").astype("string")
                 )
             except Exception:
                 pass

{upgini-1.1.316a3.dist-info → upgini-1.1.316a5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.1.316a3
+Version: 1.1.316a5
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/

{upgini-1.1.316a3.dist-info → upgini-1.1.316a5.dist-info}/RECORD

@@ -1,10 +1,10 @@
-upgini/__about__.py,sha256=xP_PhI7jmSCABPEedhQOlt9k8Njn3IHiI7PyPcsXGQQ,26
+upgini/__about__.py,sha256=5SaWm460mZelKwFqDXwqvCt7MFnWhCTJ17oJESSgrVA,26
 upgini/__init__.py,sha256=Xs0YFVBu1KUdtZzbStGRPQtLt3YLzJnjx5nIUBlX8BE,415
 upgini/ads.py,sha256=nvuRxRx5MHDMgPr9SiU-fsqRdFaBv8p4_v1oqiysKpc,2714
-upgini/dataset.py,sha256=yAWIygHejxdKXOA4g3QjtCu0VRa9at-4nPPuugCr77U,30857
+upgini/dataset.py,sha256=olZ-OHSfBNoBSCo7R5t7uCLukI2nO7afpx_A-HCiJLk,31067
 upgini/errors.py,sha256=2b_Wbo0OYhLUbrZqdLIx5jBnAsiD1Mcenh-VjR4HCTw,950
-upgini/features_enricher.py,sha256=_d8ya5RRoYN0o6mV6gda-bLdOngQ4rb1SA51SlM_TG0,188002
-upgini/http.py,sha256=gCN5ru_I6JNHk-m6-Ckjhd23iMzOAzDSLb0tSEcxkC4,43068
+upgini/features_enricher.py,sha256=9l8C3p6OaLkgE9O_kln_uJxqY1A7qqDgi5_l7X7ukeE,188002
+upgini/http.py,sha256=21asexflvavydzCOONJDGQBtQanCElrbnqLXakJ9Cu8,42880
 upgini/lazy_import.py,sha256=74gQ8JuA48BGRLxAo7lNHNKY2D2emMxrUxKGdxVGhuY,1012
 upgini/metadata.py,sha256=osmzdNESeh7yP3BZday6N9Q3eaIHfzhhRM1d6NSgcf0,11223
 upgini/metrics.py,sha256=Tu5cN8RlhOSSMWUTXRSkdl8SWBqR1N_2eJpBum9pZxc,30926
@@ -16,8 +16,8 @@ upgini/ads_management/ads_manager.py,sha256=igVbN2jz80Umb2BUJixmJVj-zx8unoKpecVo
 upgini/autofe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 upgini/autofe/all_operands.py,sha256=3LiH9iU-ArGmYpS8FHWH7yCFx40ILfvlSXJlKIa75BQ,2542
 upgini/autofe/binary.py,sha256=xRBT7RNqQ7pprz6cRpO1KnvZCb7PvU3QXBfaP6Omqi4,7425
-upgini/autofe/date.py,sha256=aKuEsguYSrFdFiLd6tBLVH4TiQ3JFMo_49_Ajp8eKQg,9208
-upgini/autofe/feature.py,sha256=CivPkE7YrAtDrgF8WhVPnDAnNDR8gbRQ-8_hXiQE6ew,14234
+upgini/autofe/date.py,sha256=ku3kcmzpPmyUmpXHIBwT6JCIkaslRknW8DifUXvFnG8,9762
+upgini/autofe/feature.py,sha256=gwGWY2UcX_0wHAvfEiu1rRU7GFZyzMWZIaPVcf6kD80,14223
 upgini/autofe/groupby.py,sha256=r-xl_keZZgm_tpiEoDhjYSkT6NHv7a4cRQR4wJ4uCp8,3263
 upgini/autofe/operand.py,sha256=uk883RaNqgXqtkaRqA1re1d9OFnnpv0JVvelYx09Yw0,2943
 upgini/autofe/unary.py,sha256=RiK-Fz3fgjPlqWWfro6x7qChjEZ8W8RTnl5-MT1kaQA,4218
@@ -51,13 +51,13 @@ upgini/utils/features_validator.py,sha256=PgKNt5dyqfErTvjtRNNUS9g7GFqHBtAtnsfA-V
 upgini/utils/format.py,sha256=Yv5cvvSs2bOLUzzNu96Pu33VMDNbabio92QepUj41jU,243
 upgini/utils/ip_utils.py,sha256=ZZj_uQFTHhagzt-MRew__ZBOp2DdnkMrachS7PElkSE,5143
 upgini/utils/phone_utils.py,sha256=IrbztLuOJBiePqqxllfABWfYlfAjYevPhXKipl95wUI,10432
-upgini/utils/postal_code_utils.py,sha256=C899tJS8qM_ps4I3g-Ve6qzIa22O_UqwNmGFoyy9sO8,1716
+upgini/utils/postal_code_utils.py,sha256=5M0sUqH2DAr33kARWCTXR-ACyzWbjDq_-0mmEml6ZcU,1716
 upgini/utils/progress_bar.py,sha256=N-Sfdah2Hg8lXP_fV9EfUTXz_PyRt4lo9fAHoUDOoLc,1550
 upgini/utils/sklearn_ext.py,sha256=13jQS_k7v0aUtudXV6nGUEWjttPQzAW9AFYL5wgEz9k,44511
 upgini/utils/target_utils.py,sha256=BVtDmrmFMKerSUWaNOIEdzsYHIFiODdpnWbE50QDPDc,7864
 upgini/utils/track_info.py,sha256=G5Lu1xxakg2_TQjKZk4b5SvrHsATTXNVV3NbvWtT8k8,5663
 upgini/utils/warning_counter.py,sha256=dIWBB4dI5XRRJZudvIlqlIYKEiwLLPcXarsZuYRt338,227
-upgini-1.1.316a3.dist-info/METADATA,sha256=wqF_a0Mo2hFvIHf5cxVPquLOnkz0LHeIOmTdRUP7R9M,48232
-upgini-1.1.316a3.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-upgini-1.1.316a3.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
-upgini-1.1.316a3.dist-info/RECORD,,
+upgini-1.1.316a5.dist-info/METADATA,sha256=xGm11UrAxkdD9Fi3SYyek-IDOvUcDxA68Dy8cH0gQ3c,48232
+upgini-1.1.316a5.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+upgini-1.1.316a5.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
+upgini-1.1.316a5.dist-info/RECORD,,