upgini 1.2.106a3956.dev1__py3-none-any.whl → 1.2.108__py3-none-any.whl

upgini/__about__.py

@@ -1 +1 @@
-__version__ = "1.2.106a3956.dev1"
+__version__ = "1.2.108"

upgini/features_enricher.py

@@ -208,7 +208,7 @@ class FeaturesEnricher(TransformerMixin):
         self,
         search_keys: Optional[Dict[str, SearchKey]] = None,
         country_code: Optional[str] = None,
-        model_task_type: Optional[ModelTaskType] = None,
+        model_task_type: Optional[Union[ModelTaskType, str]] = None,
         api_key: Optional[str] = None,
         endpoint: Optional[str] = None,
         search_id: Optional[str] = None,
@@ -234,6 +234,7 @@ class FeaturesEnricher(TransformerMixin):
         id_columns: Optional[List[str]] = None,
         generate_search_key_features: bool = True,
         sample_config: Optional[SampleConfig] = None,
+        print_trace_id: bool = False,
         **kwargs,
     ):
         self.bundle = get_custom_bundle(custom_bundle_config)
@@ -282,6 +283,8 @@ class FeaturesEnricher(TransformerMixin):
         self.country_code = country_code
         self.__validate_search_keys(search_keys, search_id)
 
+        if model_task_type is not None:
+            self.model_task_type = ModelTaskType.parse(model_task_type)
         self.model_task_type = model_task_type
         self.endpoint = endpoint
         self._search_task: Optional[SearchTask] = None
@@ -303,6 +306,8 @@ class FeaturesEnricher(TransformerMixin):
 
             print(self.bundle.get("search_by_task_id_start"))
             trace_id = str(uuid.uuid4())
+            if self.print_trace_id:
+                print(f"@trace_id:{trace_id}")
             with MDC(trace_id=trace_id):
                 try:
                     self.logger.debug(f"FeaturesEnricher created from existing search: {search_id}")
@@ -366,6 +371,7 @@ class FeaturesEnricher(TransformerMixin):
         self.data_sources_display_handle = None
         self.autofe_features_display_handle = None
         self.report_button_handle = None
+        self.print_trace_id = print_trace_id
 
     def _get_sample_config(self, sample_config: Optional[SampleConfig] = None):
         sample_config = sample_config or SampleConfig(force_sample_size=Dataset.FORCE_SAMPLE_SIZE)
@@ -461,6 +467,8 @@ class FeaturesEnricher(TransformerMixin):
             Otherwise, return all features from input and only selected features from data sources.
         """
         trace_id = str(uuid.uuid4())
+        if self.print_trace_id:
+            print(f"@trace_id:{trace_id}")
         start_time = time.time()
         auto_fe_parameters = AutoFEParameters() if auto_fe_parameters is None else auto_fe_parameters
         search_progress = SearchProgress(0.0, ProgressStage.START_FIT)
@@ -619,6 +627,8 @@ class FeaturesEnricher(TransformerMixin):
         self.warning_counter.reset()
         auto_fe_parameters = AutoFEParameters() if auto_fe_parameters is None else auto_fe_parameters
         trace_id = str(uuid.uuid4())
+        if self.print_trace_id:
+            print(f"@trace_id:{trace_id}")
         start_time = time.time()
         with MDC(trace_id=trace_id):
             if len(args) > 0:
@@ -4633,65 +4643,78 @@ if response.status_code == 200:
         eval_set: Union[Tuple, None] = None,
     ):
         def dump_task(X_, y_, eval_set_):
-            try:
-                if isinstance(X_, pd.Series):
-                    X_ = X_.to_frame()
-
-                with tempfile.TemporaryDirectory() as tmp_dir:
-                    X_.to_parquet(f"{tmp_dir}/x.parquet", compression="zstd")
-                    x_digest_sha256 = self.rest_client.compute_file_digest(f"{tmp_dir}/x.parquet")
-                    if self.rest_client.is_file_uploaded(trace_id, x_digest_sha256):
-                        self.logger.info(f"File x.parquet was already uploaded with digest {x_digest_sha256}, skipping")
-                    else:
-                        self.rest_client.dump_input_file(trace_id, f"{tmp_dir}/x.parquet", "x.parquet")
-
-                    if y_ is not None:
-                        if isinstance(y_, pd.Series):
-                            y_ = y_.to_frame()
-                        y_.to_parquet(f"{tmp_dir}/y.parquet", compression="zstd")
-                        y_digest_sha256 = self.rest_client.compute_file_digest(f"{tmp_dir}/y.parquet")
-                        if self.rest_client.is_file_uploaded(trace_id, y_digest_sha256):
+            with MDC(trace_id=trace_id):
+                try:
+                    if isinstance(X_, pd.Series):
+                        X_ = X_.to_frame()
+
+                    with tempfile.TemporaryDirectory() as tmp_dir:
+                        X_.to_parquet(f"{tmp_dir}/x.parquet", compression="zstd")
+                        x_digest_sha256 = self.rest_client.compute_file_digest(f"{tmp_dir}/x.parquet")
+                        if self.rest_client.is_file_uploaded(trace_id, x_digest_sha256):
                             self.logger.info(
-                                f"File y.parquet was already uploaded with digest {y_digest_sha256}, skipping"
+                                f"File x.parquet was already uploaded with digest {x_digest_sha256}, skipping"
                             )
                         else:
-                            self.rest_client.dump_input_file(trace_id, f"{tmp_dir}/y.parquet", "y.parquet")
-
-                        if eval_set_ is not None and len(eval_set_) > 0:
-                            for idx, (eval_x_, eval_y_) in enumerate(eval_set_):
-                                if isinstance(eval_x_, pd.Series):
-                                    eval_x_ = eval_x_.to_frame()
-                                eval_x_.to_parquet(f"{tmp_dir}/eval_x_{idx}.parquet", compression="zstd")
-                                eval_x_digest_sha256 = self.rest_client.compute_file_digest(
-                                    f"{tmp_dir}/eval_x_{idx}.parquet"
-                                )
-                                if self.rest_client.is_file_uploaded(trace_id, eval_x_digest_sha256):
-                                    self.logger.info(
-                                        f"File eval_x_{idx}.parquet was already uploaded with"
-                                        f" digest {eval_x_digest_sha256}, skipping"
-                                    )
-                                else:
-                                    self.rest_client.dump_input_file(
-                                        trace_id, f"{tmp_dir}/eval_x_{idx}.parquet", f"eval_x_{idx}.parquet"
-                                    )
+                            self.rest_client.dump_input_file(
+                                trace_id, f"{tmp_dir}/x.parquet", "x.parquet", x_digest_sha256
+                            )
 
-                                if isinstance(eval_y_, pd.Series):
-                                    eval_y_ = eval_y_.to_frame()
-                                eval_y_.to_parquet(f"{tmp_dir}/eval_y_{idx}.parquet", compression="zstd")
-                                eval_y_digest_sha256 = self.rest_client.compute_file_digest(
-                                    f"{tmp_dir}/eval_y_{idx}.parquet"
+                        if y_ is not None:
+                            if isinstance(y_, pd.Series):
+                                y_ = y_.to_frame()
+                            y_.to_parquet(f"{tmp_dir}/y.parquet", compression="zstd")
+                            y_digest_sha256 = self.rest_client.compute_file_digest(f"{tmp_dir}/y.parquet")
+                            if self.rest_client.is_file_uploaded(trace_id, y_digest_sha256):
+                                self.logger.info(
+                                    f"File y.parquet was already uploaded with digest {y_digest_sha256}, skipping"
                                 )
-                                if self.rest_client.is_file_uploaded(trace_id, eval_y_digest_sha256):
-                                    self.logger.info(
-                                        f"File eval_y_{idx}.parquet was already uploaded"
-                                        f" with digest {eval_y_digest_sha256}, skipping"
+                            else:
+                                self.rest_client.dump_input_file(
+                                    trace_id, f"{tmp_dir}/y.parquet", "y.parquet", y_digest_sha256
+                                )
+
+                            if eval_set_ is not None and len(eval_set_) > 0:
+                                for idx, (eval_x_, eval_y_) in enumerate(eval_set_):
+                                    if isinstance(eval_x_, pd.Series):
+                                        eval_x_ = eval_x_.to_frame()
+                                    eval_x_.to_parquet(f"{tmp_dir}/eval_x_{idx}.parquet", compression="zstd")
+                                    eval_x_digest_sha256 = self.rest_client.compute_file_digest(
+                                        f"{tmp_dir}/eval_x_{idx}.parquet"
                                     )
-                                else:
-                                    self.rest_client.dump_input_file(
-                                        trace_id, f"{tmp_dir}/eval_y_{idx}.parquet", f"eval_y_{idx}.parquet"
+                                    if self.rest_client.is_file_uploaded(trace_id, eval_x_digest_sha256):
+                                        self.logger.info(
+                                            f"File eval_x_{idx}.parquet was already uploaded with"
+                                            f" digest {eval_x_digest_sha256}, skipping"
+                                        )
+                                    else:
+                                        self.rest_client.dump_input_file(
+                                            trace_id,
+                                            f"{tmp_dir}/eval_x_{idx}.parquet",
+                                            f"eval_x_{idx}.parquet",
+                                            eval_x_digest_sha256,
+                                        )
+
+                                    if isinstance(eval_y_, pd.Series):
+                                        eval_y_ = eval_y_.to_frame()
+                                    eval_y_.to_parquet(f"{tmp_dir}/eval_y_{idx}.parquet", compression="zstd")
+                                    eval_y_digest_sha256 = self.rest_client.compute_file_digest(
+                                        f"{tmp_dir}/eval_y_{idx}.parquet"
                                     )
-            except Exception:
-                self.logger.warning("Failed to dump input files", exc_info=True)
+                                    if self.rest_client.is_file_uploaded(trace_id, eval_y_digest_sha256):
+                                        self.logger.info(
+                                            f"File eval_y_{idx}.parquet was already uploaded"
+                                            f" with digest {eval_y_digest_sha256}, skipping"
+                                        )
+                                    else:
+                                        self.rest_client.dump_input_file(
+                                            trace_id,
+                                            f"{tmp_dir}/eval_y_{idx}.parquet",
+                                            f"eval_y_{idx}.parquet",
+                                            eval_y_digest_sha256,
+                                        )
+                except Exception:
+                    self.logger.warning("Failed to dump input files", exc_info=True)
 
         try:
             Thread(target=dump_task, args=(X, y, eval_set), daemon=True).start()

upgini/http.py

@@ -274,7 +274,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"
+    SEARCH_DUMP_INPUT_FILE_FMT = SERVICE_ROOT_V2 + "search/dump-input-file?digest={0}"
     TRANSFORM_USAGE_FMT = SERVICE_ROOT_V2 + "user/transform-usage"
 
     UPLOAD_USER_ADS_URI = SERVICE_ROOT + "ads/upload"
@@ -406,8 +406,8 @@ class _RestClient:
         meaning_types = [_RestClient.meaning_type_by_name(name, metadata) for name in search_key_names]
         return [meaning_type.value for meaning_type in meaning_types if meaning_type is not None]
 
-    def dump_input_file(self, trace_id: str, path: str, file_name: str):
-        api_path = self.SEARCH_DUMP_INPUT_FILE_FMT
+    def dump_input_file(self, trace_id: str, path: str, file_name: str, digest: str):
+        api_path = self.SEARCH_DUMP_INPUT_FILE_FMT.format(digest)
         with open(path, "rb") as file:
             files = {"file": (file_name, file, "application/octet-stream")}
             self._with_unauth_retry(

upgini/metadata.py

@@ -162,6 +162,15 @@ class ModelTaskType(Enum):
     def is_classification(self) -> bool:
         return self in [ModelTaskType.BINARY, ModelTaskType.MULTICLASS]
 
+    @staticmethod
+    def parse(task_type: Any) -> "ModelTaskType":
+        if isinstance(task_type, ModelTaskType):
+            return task_type
+        elif isinstance(task_type, str):
+            return ModelTaskType(task_type.upper())
+        else:
+            raise ValueError(f"Invalid task type: {task_type}")
+
 
 class ModelLabelType(Enum):
     GINI = "gini"

{upgini-1.2.106a3956.dev1.dist-info → upgini-1.2.108.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.106a3956.dev1
+Version: 1.2.108
 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.2.106a3956.dev1.dist-info → upgini-1.2.108.dist-info}/RECORD

@@ -1,11 +1,11 @@
-upgini/__about__.py,sha256=wMowjQ4NPzJq0VApX8igfNH42M4gS9xSnMJkWS4GHwk,34
+upgini/__about__.py,sha256=MtpgyPilS-p0uCXLJRENxbcYk2BQX6y8kTyPV7OfGCU,24
 upgini/__init__.py,sha256=LXSfTNU0HnlOkE69VCxkgIKDhWP-JFo_eBQ71OxTr5Y,261
 upgini/ads.py,sha256=nvuRxRx5MHDMgPr9SiU-fsqRdFaBv8p4_v1oqiysKpc,2714
 upgini/dataset.py,sha256=e6JDYTZ2AwC5aF-dqclKZKkiKrHo2f6cFmMQO2ZZmjM,32724
 upgini/errors.py,sha256=2b_Wbo0OYhLUbrZqdLIx5jBnAsiD1Mcenh-VjR4HCTw,950
-upgini/features_enricher.py,sha256=dBCBqAhzVHqRB2b1sPy9wzSi9XtIzeb6uArmJhcjj_8,220370
-upgini/http.py,sha256=DNcoS7qdxG0mOJn6I8r6O5I6XdIJTdzDzW3hkz3NgG4,45443
-upgini/metadata.py,sha256=vsbbHyPCP3Rs8WkeDgQg99uAA_zmsbDStAT-NwDYhO4,12455
+upgini/features_enricher.py,sha256=BHZcpkUl7ncSdTgiMxk_a6oD1pLEkKPOm5wxebK8TsU,221609
+upgini/http.py,sha256=zeAZvT6IAzOs9jQ3WG8mJBANLajgvv2LZePFzKz004w,45482
+upgini/metadata.py,sha256=9_0lFEWPpIHRBW-xWYSEcwPzICTC6_bQ6dUUlE75Xns,12773
 upgini/metrics.py,sha256=V2SP6NS5bfFHzRqufeKVsCXME1yG4t_8Dmk2E3zKdYk,45715
 upgini/search_task.py,sha256=Q5HjBpLIB3OCxAD1zNv5yQ3ZNJx696WCK_-H35_y7Rs,17912
 upgini/spinner.py,sha256=4iMd-eIe_BnkqFEMIliULTbj6rNI2HkN_VJ4qYe0cUc,1118
@@ -71,7 +71,7 @@ upgini/utils/target_utils.py,sha256=i3Xt5l9ybB2_nF_ma5cfPuL3OeFTs2dY2xDI0p4Azpg,
 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.106a3956.dev1.dist-info/METADATA,sha256=8gm0u1avCVEV1kBGkO9qYEFZHuRJ-lv_c4i0CfyO1Hg,49539
-upgini-1.2.106a3956.dev1.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-upgini-1.2.106a3956.dev1.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
-upgini-1.2.106a3956.dev1.dist-info/RECORD,,
+upgini-1.2.108.dist-info/METADATA,sha256=O9XA7uFUs-bvKrspv0HZi6vD_v8ZLSVckmDW5N5CeRo,49529
+upgini-1.2.108.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+upgini-1.2.108.dist-info/licenses/LICENSE,sha256=5RRzgvdJUu3BUDfv4bzVU6FqKgwHlIay63pPCSmSgzw,1514
+upgini-1.2.108.dist-info/RECORD,,

{upgini-1.2.106a3956.dev1.dist-info → upgini-1.2.108.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