snowflake-ml-python 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

snowflake/ml/modeling/model_selection/grid_search_cv.py

@@ -20,6 +20,7 @@ from snowflake.ml.model.model_signature import (
     FeatureSpec,
     ModelSignature,
     _infer_signature,
+    _rename_signature_with_snowflake_identifiers,
 )
 from snowflake.ml.modeling._internal.estimator_utils import (
     gather_dependencies,
@@ -330,12 +331,15 @@ class GridSearchCV(BaseTransformer):
         )
         self._sklearn_object = model_trainer.train()
         self._is_fitted = True
-        self._get_model_signatures(dataset)
+        self._generate_model_signatures(dataset)
         return self
 
-    def _batch_inference_validate_snowpark(self, dataset: DataFrame, inference_method: str) -> List[str]:
-        """Util method to run validate that batch inference can be run on a snowpark dataframe and
-        return the available package that exists in the snowflake anaconda channel
+    def _batch_inference_validate_snowpark(
+        self,
+        dataset: DataFrame,
+        inference_method: str,
+    ) -> None:
+        """Util method to run validate that batch inference can be run on a snowpark dataframe.
 
         Args:
             dataset: snowpark dataframe
@@ -345,8 +349,6 @@ class GridSearchCV(BaseTransformer):
             SnowflakeMLException: If the estimator is not fitted, raise error
             SnowflakeMLException: If the session is None, raise error
 
-        Returns:
-            A list of available package that exists in the snowflake anaconda channel
         """
         if not self._is_fitted:
             raise exceptions.SnowflakeMLException(
@@ -362,10 +364,6 @@ class GridSearchCV(BaseTransformer):
                 error_code=error_codes.NOT_FOUND,
                 original_exception=ValueError("Session must not specified for snowpark dataset."),
             )
-        # Validate that key package version in user workspace are supported in snowflake conda channel
-        return pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
-            pkg_versions=self._get_dependencies(), session=session, subproject=_SUBPROJECT
-        )
 
     @available_if(original_estimator_has_callable("predict"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -384,6 +382,9 @@ class GridSearchCV(BaseTransformer):
 
         Returns:
             Transformed dataset.
+
+        Raises:
+            SnowflakeMLException: when the output column(s) doesn't exist in the model signature, raise error
         """
         super()._check_dataset_type(dataset)
 
@@ -396,13 +397,23 @@ class GridSearchCV(BaseTransformer):
             expected_type_inferred = ""
             # infer the datatype from label columns
             if "predict" in self.model_signatures:
-                expected_type_inferred = convert_sp_to_sf_type(
-                    self.model_signatures["predict"].outputs[0].as_snowpark_type()
-                )
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
+                # Batch inference takes a single expected output column type. Use the first columns type for now.
+                label_cols_signatures = [
+                    row for row in self.model_signatures["predict"].outputs if row.name in self.output_cols
+                ]
+                if len(label_cols_signatures) == 0:
+                    error_str = (
+                        f"Output columns {self.output_cols} do not match"
+                        f"model signatures {self.model_signatures['predict'].outputs}."
+                    )
+                    raise exceptions.SnowflakeMLException(
+                        error_code=error_codes.INVALID_ATTRIBUTE,
+                        original_exception=ValueError(error_str),
+                    )
+
+                expected_type_inferred = convert_sp_to_sf_type(label_cols_signatures[0].as_snowpark_type())
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
 
             assert isinstance(
                 dataset._session, Session
@@ -460,7 +471,8 @@ class GridSearchCV(BaseTransformer):
         inference_method = "transform"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -519,7 +531,8 @@ class GridSearchCV(BaseTransformer):
         inference_method = "predict_proba"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -579,7 +592,8 @@ class GridSearchCV(BaseTransformer):
         inference_method = "predict_log_proba"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -639,7 +653,8 @@ class GridSearchCV(BaseTransformer):
         inference_method = "decision_function"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -700,7 +715,8 @@ class GridSearchCV(BaseTransformer):
         transform_kwargs: BatchInferenceKwargsTypedDict = dict()
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -751,17 +767,15 @@ class GridSearchCV(BaseTransformer):
         transform_kwargs: ScoreKwargsTypedDict = dict()
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method="score",
-            )
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method="score")
+            self._deps = self._get_dependencies()
             selected_cols = self._get_active_columns()
             if len(selected_cols) > 0:
                 dataset = dataset.select(selected_cols)
             assert isinstance(dataset._session, Session)  # keep mypy happy
             transform_kwargs = dict(
                 session=dataset._session,
-                dependencies=["snowflake-snowpark-python"] + self._deps,
+                dependencies=self._deps,
                 score_sproc_imports=["sklearn"],
             )
         elif isinstance(dataset, pd.DataFrame):
@@ -785,12 +799,22 @@ class GridSearchCV(BaseTransformer):
 
         return output_score
 
-    def _get_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
+    def to_sklearn(self) -> sklearn.model_selection.GridSearchCV:
+        """
+        Get sklearn.model_selection.GridSearchCV object.
+        """
+        assert self._sklearn_object is not None
+        return self._sklearn_object
+
+    def _get_dependencies(self) -> List[str]:
+        return self._deps
+
+    def _generate_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
         self._model_signature_dict = dict()
 
         PROB_FUNCTIONS = ["predict_log_proba", "predict_proba", "decision_function"]
 
-        inputs = list(_infer_signature(dataset[self.input_cols], "input"))
+        inputs = list(_infer_signature(dataset[self.input_cols], "input", use_snowflake_identifiers=True))
         outputs: List[BaseFeatureSpec] = []
         if hasattr(self, "predict"):
             # keep mypy happy
@@ -798,18 +822,20 @@ class GridSearchCV(BaseTransformer):
             # For classifier, the type of predict is the same as the type of label
             if self._sklearn_object._estimator_type == "classifier":
                 # label columns is the desired type for output
-                outputs = list(_infer_signature(dataset[self.label_cols], "output"))
+                outputs = list(_infer_signature(dataset[self.label_cols], "output", use_snowflake_identifiers=True))
                 # rename the output columns
                 outputs = list(model_signature_utils.rename_features(outputs, self.output_cols))
                 self._model_signature_dict["predict"] = ModelSignature(
                     inputs, ([] if self._drop_input_cols else inputs) + outputs
                 )
+
             # For regressor, the type of predict is float64
             elif self._sklearn_object._estimator_type == "regressor":
                 outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in self.output_cols]
                 self._model_signature_dict["predict"] = ModelSignature(
                     inputs, ([] if self._drop_input_cols else inputs) + outputs
                 )
+
         for prob_func in PROB_FUNCTIONS:
             if hasattr(self, prob_func):
                 output_cols_prefix: str = f"{prob_func}_"
@@ -819,6 +845,12 @@ class GridSearchCV(BaseTransformer):
                     inputs, ([] if self._drop_input_cols else inputs) + outputs
                 )
 
+        # Output signature names may still need to be renamed, since they were not created with `_infer_signature`.
+        items = list(self._model_signature_dict.items())
+        for method, signature in items:
+            signature._outputs = _rename_signature_with_snowflake_identifiers(signature._outputs)
+            self._model_signature_dict[method] = signature
+
     @property
     def model_signatures(self) -> Dict[str, ModelSignature]:
         """Returns model signature of current class.
@@ -827,7 +859,7 @@ class GridSearchCV(BaseTransformer):
             SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
 
         Returns:
-            Dict[str, ModelSignature]: each method and its input output signature
+            each method and its input output signature
         """
         if self._model_signature_dict is None:
             raise exceptions.SnowflakeMLException(
@@ -835,13 +867,3 @@ class GridSearchCV(BaseTransformer):
                 original_exception=RuntimeError("Estimator not fitted before accessing property model_signatures!"),
             )
         return self._model_signature_dict
-
-    def to_sklearn(self) -> sklearn.model_selection.GridSearchCV:
-        """
-        Get sklearn.model_selection.GridSearchCV object.
-        """
-        assert self._sklearn_object is not None
-        return self._sklearn_object
-
-    def _get_dependencies(self) -> List[str]:
-        return self._deps

snowflake/ml/modeling/model_selection/randomized_search_cv.py

@@ -17,6 +17,7 @@ from snowflake.ml.model.model_signature import (
     FeatureSpec,
     ModelSignature,
     _infer_signature,
+    _rename_signature_with_snowflake_identifiers,
 )
 from snowflake.ml.modeling._internal.estimator_utils import (
     gather_dependencies,
@@ -343,11 +344,25 @@ class RandomizedSearchCV(BaseTransformer):
         )
         self._sklearn_object = model_trainer.train()
         self._is_fitted = True
-        self._get_model_signatures(dataset)
+        self._generate_model_signatures(dataset)
         return self
 
-    def _batch_inference_validate_snowpark(self, dataset: DataFrame, inference_method: str) -> List[str]:
-        """Util method to run validate that batch inference can be run on a snowpark dataframe."""
+    def _batch_inference_validate_snowpark(
+        self,
+        dataset: DataFrame,
+        inference_method: str,
+    ) -> None:
+        """Util method to run validate that batch inference can be run on a snowpark dataframe.
+
+        Args:
+            dataset: snowpark dataframe
+            inference_method: the inference method such as predict, score...
+
+        Raises:
+            SnowflakeMLException: If the estimator is not fitted, raise error
+            SnowflakeMLException: If the session is None, raise error
+
+        """
         if not self._is_fitted:
             raise exceptions.SnowflakeMLException(
                 error_code=error_codes.METHOD_NOT_ALLOWED,
@@ -362,10 +377,6 @@ class RandomizedSearchCV(BaseTransformer):
                 error_code=error_codes.NOT_FOUND,
                 original_exception=ValueError("Session must not specified for snowpark dataset."),
             )
-        # Validate that key package version in user workspace are supported in snowflake conda channel
-        return pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
-            pkg_versions=self._get_dependencies(), session=session, subproject=_SUBPROJECT
-        )
 
     @available_if(original_estimator_has_callable("predict"))  # type: ignore[misc]
     @telemetry.send_api_usage_telemetry(
@@ -383,6 +394,9 @@ class RandomizedSearchCV(BaseTransformer):
 
         Returns:
             Transformed dataset.
+
+        Raises:
+            SnowflakeMLException: when the output column(s) doesn't exist in the model signature, raise error
         """
         super()._check_dataset_type(dataset)
 
@@ -395,13 +409,24 @@ class RandomizedSearchCV(BaseTransformer):
             expected_type_inferred = ""
             # infer the datatype from label columns
             if "predict" in self.model_signatures:
-                expected_type_inferred = convert_sp_to_sf_type(
-                    self.model_signatures["predict"].outputs[0].as_snowpark_type()
-                )
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method=inference_method,
-            )
+                # Batch inference takes a single expected output column type. Use the first columns type for now.
+                label_cols_signatures = [
+                    row for row in self.model_signatures["predict"].outputs if row.name in self.output_cols
+                ]
+                if len(label_cols_signatures) == 0:
+                    error_str = (
+                        f"Output columns {self.output_cols} do not match"
+                        f"model signatures {self.model_signatures['predict'].outputs}."
+                    )
+                    raise exceptions.SnowflakeMLException(
+                        error_code=error_codes.INVALID_ATTRIBUTE,
+                        original_exception=ValueError(error_str),
+                    )
+
+                expected_type_inferred = convert_sp_to_sf_type(label_cols_signatures[0].as_snowpark_type())
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -457,7 +482,9 @@ class RandomizedSearchCV(BaseTransformer):
         inference_method = "transform"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -515,7 +542,9 @@ class RandomizedSearchCV(BaseTransformer):
         inference_method = "predict_proba"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -575,7 +604,9 @@ class RandomizedSearchCV(BaseTransformer):
         inference_method = "predict_log_proba"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -634,7 +665,9 @@ class RandomizedSearchCV(BaseTransformer):
         inference_method = "decision_function"
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -695,7 +728,9 @@ class RandomizedSearchCV(BaseTransformer):
         transform_kwargs: BatchInferenceKwargsTypedDict = dict()
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method=inference_method)
+            self._deps = self._get_dependencies()
+
             assert isinstance(
                 dataset._session, Session
             )  # mypy does not recognize the check in _batch_inference_validate_snowpark()
@@ -745,10 +780,9 @@ class RandomizedSearchCV(BaseTransformer):
         transform_kwargs: ScoreKwargsTypedDict = dict()
 
         if isinstance(dataset, DataFrame):
-            self._deps = self._batch_inference_validate_snowpark(
-                dataset=dataset,
-                inference_method="score",
-            )
+            self._batch_inference_validate_snowpark(dataset=dataset, inference_method="score")
+            self._deps = self._get_dependencies()
+
             selected_cols = self._get_active_columns()
             if len(selected_cols) > 0:
                 dataset = dataset.select(selected_cols)
@@ -756,7 +790,7 @@ class RandomizedSearchCV(BaseTransformer):
             assert isinstance(dataset._session, Session)  # keep mypy happy
             transform_kwargs = dict(
                 session=dataset._session,
-                dependencies=["snowflake-snowpark-python"] + self._deps,
+                dependencies=self._deps,
                 score_sproc_imports=["sklearn"],
             )
         elif isinstance(dataset, pd.DataFrame):
@@ -780,12 +814,22 @@ class RandomizedSearchCV(BaseTransformer):
 
         return output_score
 
-    def _get_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
+    def to_sklearn(self) -> sklearn.model_selection.RandomizedSearchCV:
+        """
+        Get sklearn.model_selection.RandomizedSearchCV object.
+        """
+        assert self._sklearn_object is not None
+        return self._sklearn_object
+
+    def _get_dependencies(self) -> List[str]:
+        return self._deps
+
+    def _generate_model_signatures(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
         self._model_signature_dict = dict()
 
         PROB_FUNCTIONS = ["predict_log_proba", "predict_proba", "decision_function"]
 
-        inputs = list(_infer_signature(dataset[self.input_cols], "input"))
+        inputs = list(_infer_signature(dataset[self.input_cols], "input", use_snowflake_identifiers=True))
         outputs: List[BaseFeatureSpec] = []
         if hasattr(self, "predict"):
             # keep mypy happy
@@ -793,18 +837,20 @@ class RandomizedSearchCV(BaseTransformer):
             # For classifier, the type of predict is the same as the type of label
             if self._sklearn_object._estimator_type == "classifier":
                 # label columns is the desired type for output
-                outputs = list(_infer_signature(dataset[self.label_cols], "output"))
+                outputs = list(_infer_signature(dataset[self.label_cols], "output", use_snowflake_identifiers=True))
                 # rename the output columns
                 outputs = list(model_signature_utils.rename_features(outputs, self.output_cols))
                 self._model_signature_dict["predict"] = ModelSignature(
                     inputs, ([] if self._drop_input_cols else inputs) + outputs
                 )
+
             # For regressor, the type of predict is float64
             elif self._sklearn_object._estimator_type == "regressor":
                 outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in self.output_cols]
                 self._model_signature_dict["predict"] = ModelSignature(
                     inputs, ([] if self._drop_input_cols else inputs) + outputs
                 )
+
         for prob_func in PROB_FUNCTIONS:
             if hasattr(self, prob_func):
                 output_cols_prefix: str = f"{prob_func}_"
@@ -814,6 +860,12 @@ class RandomizedSearchCV(BaseTransformer):
                     inputs, ([] if self._drop_input_cols else inputs) + outputs
                 )
 
+        # Output signature names may still need to be renamed, since they were not created with `_infer_signature`.
+        items = list(self._model_signature_dict.items())
+        for method, signature in items:
+            signature._outputs = _rename_signature_with_snowflake_identifiers(signature._outputs)
+            self._model_signature_dict[method] = signature
+
     @property
     def model_signatures(self) -> Dict[str, ModelSignature]:
         """Returns model signature of current class.
@@ -822,7 +874,7 @@ class RandomizedSearchCV(BaseTransformer):
             SnowflakeMLException: If estimator is not fitted, then model signature cannot be inferred
 
         Returns:
-            Dict[str, ModelSignature]: each method and its input output signature
+            each method and its input output signature
         """
         if self._model_signature_dict is None:
             raise exceptions.SnowflakeMLException(
@@ -830,13 +882,3 @@ class RandomizedSearchCV(BaseTransformer):
                 original_exception=RuntimeError("Estimator not fitted before accessing property model_signatures!"),
             )
         return self._model_signature_dict
-
-    def to_sklearn(self) -> sklearn.model_selection.RandomizedSearchCV:
-        """
-        Get sklearn.model_selection.RandomizedSearchCV object.
-        """
-        assert self._sklearn_object is not None
-        return self._sklearn_object
-
-    def _get_dependencies(self) -> List[str]:
-        return self._deps