snowflake-ml-python 1.1.1__py3-none-any.whl → 1.1.2__py3-none-any.whl

snowflake/ml/modeling/model_selection/grid_search_cv.py

@@ -2,13 +2,13 @@
 # This code is auto-generated using the sklearn_wrapper_template.py_template template.
 # Do not modify the auto-generated code(except automatic reformatting by precommit hooks).
 #
-from typing import Dict, Iterable, List, Optional, Set, Union
+from typing import Any, Dict, Iterable, List, Optional, Set, Union
 from uuid import uuid4
 
+import cloudpickle as cp
 import numpy as np
 import pandas as pd
 import sklearn.model_selection
-from sklearn.model_selection import ParameterGrid
 from sklearn.utils.metaestimators import available_if
 
 from snowflake.ml._internal import telemetry
@@ -25,13 +25,12 @@ from snowflake.ml.model.model_signature import (
 from snowflake.ml.modeling._internal.estimator_protocols import CVHandlers
 from snowflake.ml.modeling._internal.estimator_utils import (
     gather_dependencies,
-    is_single_node,
     original_estimator_has_callable,
     transform_snowml_obj_to_sklearn_obj,
     validate_sklearn_args,
 )
+from snowflake.ml.modeling._internal.model_trainer_builder import ModelTrainerBuilder
 from snowflake.ml.modeling._internal.snowpark_handlers import (
-    SklearnModelSelectionWrapperProvider,
     SnowparkHandlers as HandlersImpl,
 )
 from snowflake.ml.modeling.framework.base import BaseTransformer
@@ -53,19 +52,54 @@ class GridSearchCV(BaseTransformer):
 
     Parameters
     ----------
-    estimator : estimator object
+    estimator: estimator object
         This is assumed to implement the scikit-learn estimator interface.
         Either estimator needs to provide a ``score`` function,
         or ``scoring`` must be passed.
 
-    param_grid : dict or list of dictionaries
+    param_grid: dict or list of dictionaries
         Dictionary with parameters names (`str`) as keys and lists of
         parameter settings to try as values, or a list of such
         dictionaries, in which case the grids spanned by each dictionary
         in the list are explored. This enables searching over any sequence
         of parameter settings.
 
-    scoring : str, callable, list, tuple or dict, default=None
+    input_cols: Optional[Union[str, List[str]]]
+        A string or list of strings representing column names that contain features.
+        If this parameter is not specified, all columns in the input DataFrame except
+        the columns specified by label_cols and sample-weight_col parameters are
+        considered input columns.
+
+    label_cols: Optional[Union[str, List[str]]]
+        A string or list of strings representing column names that contain labels.
+        This is a required param for estimators, as there is no way to infer these
+        columns. If this parameter is not specified, then object is fitted without
+        labels(Like a transformer).
+
+    output_cols: Optional[Union[str, List[str]]]
+        A string or list of strings representing column names that will store the
+        output of predict and transform operations. The length of output_cols mus
+        match the expected number of output columns from the specific estimator or
+        transformer class used.
+        If this parameter is not specified, output column names are derived by
+        adding an OUTPUT_ prefix to the label column names. These inferred output
+        column names work for estimator's predict() method, but output_cols must
+        be set explicitly for transformers.
+
+    passthrough_cols: A string or a list of strings indicating column names to be excluded from any
+        operations (such as train, transform, or inference). These specified column(s)
+        will remain untouched throughout the process. This option is helpful in scenarios
+        requiring automatic input_cols inference, but need to avoid using specific
+        columns, like index columns, during training or inference.
+
+    sample_weight_col: Optional[str]
+        A string representing the column name containing the examples’ weights.
+        This argument is only required when working with weighted datasets.
+
+    drop_input_cols: Optional[bool], default=False
+        If set, the response of predict(), transform() methods will not contain input columns.
+
+    scoring: str, callable, list, tuple or dict, default=None
         Strategy to evaluate the performance of the cross-validated model on
         the test set.
 
@@ -83,13 +117,13 @@ class GridSearchCV(BaseTransformer):
 
         See :ref:`multimetric_grid_search` for an example.
 
-    n_jobs : int, default=None
+    n_jobs: int, default=None
         Number of jobs to run in parallel.
         ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
         ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
         for more details.
 
-    refit : bool, str, or callable, default=True
+    refit: bool, str, or callable, default=True
         Refit an estimator using the best found parameters on the whole
         dataset.
 
@@ -120,7 +154,7 @@ class GridSearchCV(BaseTransformer):
         to see how to design a custom selection strategy using a callable
         via `refit`.
 
-    cv : int, cross-validation generator or an iterable, default=None
+    cv: int, cross-validation generator or an iterable, default=None
         Determines the cross-validation splitting strategy.
         Possible inputs for cv are:
 
@@ -137,7 +171,7 @@ class GridSearchCV(BaseTransformer):
         Refer :ref:`User Guide <cross_validation>` for the various
         cross-validation strategies that can be used here.
 
-    verbose : int
+    verbose: int
         Controls the verbosity: the higher, the more messages.
 
         - >1 : the computation time for each fold and parameter candidate is
@@ -146,7 +180,7 @@ class GridSearchCV(BaseTransformer):
         - >3 : the fold and candidate parameter indexes are also displayed
           together with the starting time of the computation.
 
-    pre_dispatch : int, or str, default='2*n_jobs'
+    pre_dispatch: int, or str, default='2*n_jobs'
         Controls the number of jobs that get dispatched during parallel
         execution. Reducing this number can be useful to avoid an
         explosion of memory consumption when more jobs get dispatched
@@ -163,13 +197,13 @@ class GridSearchCV(BaseTransformer):
             - A str, giving an expression as a function of n_jobs,
               as in '2*n_jobs'
 
-    error_score : 'raise' or numeric, default=np.nan
+    error_score: 'raise' or numeric, default=np.nan
         Value to assign to the score if an error occurs in estimator fitting.
         If set to 'raise', the error is raised. If a numeric value is given,
         FitFailedWarning is raised. This parameter does not affect the refit
         step, which will always raise the error.
 
-    return_train_score : bool, default=False
+    return_train_score: bool, default=False
         If ``False``, the ``cv_results_`` attribute will not include training
         scores.
         Computing training scores is used to get insights on how different
@@ -177,41 +211,6 @@ class GridSearchCV(BaseTransformer):
         However computing the scores on the training set can be computationally
         expensive and is not strictly required to select the parameters that
         yield the best generalization performance.
-
-    input_cols : Optional[Union[str, List[str]]]
-        A string or list of strings representing column names that contain features.
-        If this parameter is not specified, all columns in the input DataFrame except
-        the columns specified by label_cols and sample-weight_col parameters are
-        considered input columns.
-
-    label_cols : Optional[Union[str, List[str]]]
-        A string or list of strings representing column names that contain labels.
-        This is a required param for estimators, as there is no way to infer these
-        columns. If this parameter is not specified, then object is fitted without
-        labels(Like a transformer).
-
-    output_cols: Optional[Union[str, List[str]]]
-        A string or list of strings representing column names that will store the
-        output of predict and transform operations. The length of output_cols mus
-        match the expected number of output columns from the specific estimator or
-        transformer class used.
-        If this parameter is not specified, output column names are derived by
-        adding an OUTPUT_ prefix to the label column names. These inferred output
-        column names work for estimator's predict() method, but output_cols must
-        be set explicitly for transformers.
-
-    passthrough_cols: A string or a list of strings indicating column names to be excluded from any
-        operations (such as train, transform, or inference). These specified column(s)
-        will remain untouched throughout the process. This option is helpful in scenarios
-        requiring automatic input_cols inference, but need to avoid using specific
-        columns, like index columns, during training or inference.
-
-    sample_weight_col: Optional[str]
-        A string representing the column name containing the examples’ weights.
-        This argument is only required when working with weighted datasets.
-
-    drop_input_cols: Optional[bool], default=False
-        If set, the response of predict(), transform() methods will not contain input columns.
     """
     _ENABLE_DISTRIBUTED = True
 
@@ -236,7 +235,11 @@ class GridSearchCV(BaseTransformer):
         sample_weight_col: Optional[str] = None,
     ) -> None:
         super().__init__()
-        deps: Set[str] = set(SklearnModelSelectionWrapperProvider().dependencies)
+        deps: Set[str] = {
+            f"numpy=={np.__version__}",
+            f"scikit-learn=={sklearn.__version__}",
+            f"cloudpickle=={cp.__version__}",
+        }
         deps = deps | gather_dependencies(estimator)
         self._deps = list(deps)
         estimator = transform_snowml_obj_to_sklearn_obj(estimator)
@@ -253,7 +256,7 @@ class GridSearchCV(BaseTransformer):
             "return_train_score": (return_train_score, False, False),
         }
         cleaned_up_init_args = validate_sklearn_args(args=init_args, klass=sklearn.model_selection.GridSearchCV)
-        self._sklearn_object = sklearn.model_selection.GridSearchCV(
+        self._sklearn_object: Any = sklearn.model_selection.GridSearchCV(
             **cleaned_up_init_args,
         )
         self._model_signature_dict: Optional[Dict[str, ModelSignature]] = None
@@ -266,7 +269,6 @@ class GridSearchCV(BaseTransformer):
         self._handlers: CVHandlers = HandlersImpl(
             class_name=self.__class__.__name__,
             subproject=_SUBPROJECT,
-            wrapper_provider=SklearnModelSelectionWrapperProvider(),
         )
 
     def _get_rand_id(self) -> str:
@@ -294,10 +296,6 @@ class GridSearchCV(BaseTransformer):
         For more details on this function, see [sklearn.model_selection.GridSearchCV.fit]
         (https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html#sklearn.model_selection.GridSearchCV.fit)
 
-
-        Raises:
-            TypeError: Supported dataset types: snowpark.DataFrame, pandas.DataFrame.
-
         Args:
             dataset: Union[snowflake.snowpark.DataFrame, pandas.DataFrame]
                 Snowpark or Pandas DataFrame.
@@ -306,70 +304,37 @@ class GridSearchCV(BaseTransformer):
             self
         """
         self._infer_input_output_cols(dataset)
-        if isinstance(dataset, pd.DataFrame):
-            self._estimator = self._handlers.fit_pandas(
-                dataset, self._sklearn_object, self.input_cols, self.label_cols, self.sample_weight_col
-            )
-        elif isinstance(dataset, DataFrame):
-            self._fit_snowpark(dataset)
-        else:
-            raise TypeError(
-                f"Unexpected dataset type: {type(dataset)}."
-                "Supported dataset types: snowpark.DataFrame, pandas.DataFrame."
+        if self._sklearn_object.n_jobs is None:
+            self._sklearn_object.n_jobs = -1
+        if isinstance(dataset, DataFrame):
+            session = dataset._session
+            assert session is not None  # keep mypy happy
+            # Validate that key package version in user workspace are supported in snowflake conda channel
+            # If customer doesn't have package in conda channel, replace the ones have the closest versions
+            self._deps = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
+                pkg_versions=self._get_dependencies(), session=session, subproject=_SUBPROJECT
             )
-        self._is_fitted = True
-        self._get_model_signatures(dataset)
-        return self
 
-    def _fit_snowpark(self, dataset: DataFrame) -> None:
-        session = dataset._session
-        assert session is not None  # keep mypy happy
-        # Validate that key package version in user workspace are supported in snowflake conda channel
-        # If customer doesn't have package in conda channel, replace the ones have the closest versions
-        self._deps = pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
-            pkg_versions=self._get_dependencies(), session=session, subproject=_SUBPROJECT
-        )
+            # Specify input columns so column pruning will be enforced
+            selected_cols = self._get_active_columns()
+            if len(selected_cols) > 0:
+                dataset = dataset.select(selected_cols)
 
-        selected_cols = self._get_active_columns()
-        if len(selected_cols) > 0:
-            dataset = dataset.select(selected_cols)
+            self._snowpark_cols = dataset.select(self.input_cols).columns
 
-        assert self._sklearn_object is not None
-        is_distributed = not is_single_node(session) and self._ENABLE_DISTRIBUTED is True
-        if is_distributed:
-            # Set the default value of the `n_jobs` attribute for the estimator.
-            # If minus one is set, it will not be abided by in the UDTF, so we set that to the default value as well.
-            if hasattr(self._sklearn_object.estimator, "n_jobs") and self._sklearn_object.estimator.n_jobs in [
-                None,
-                -1,
-            ]:
-                self._sklearn_object.estimator.n_jobs = DEFAULT_UDTF_NJOBS
-            self._sklearn_object = self._handlers.fit_search_snowpark(
-                param_grid=ParameterGrid(self._sklearn_object.param_grid),
-                dataset=dataset,
-                session=session,
-                estimator=self._sklearn_object,
-                dependencies=self._get_dependencies(),
-                udf_imports=["sklearn"],
-                input_cols=self.input_cols,
-                label_cols=self.label_cols,
-                sample_weight_col=self.sample_weight_col,
-            )
-        else:
-            # Fall back with stored procedure implementation
-            # set the parallel factor to default to minus one, to fully accelerate the cores in single node
-            if self._sklearn_object.n_jobs is None:
-                self._sklearn_object.n_jobs = -1
-
-            self._sklearn_object = self._handlers.fit_snowpark(
-                dataset,
-                session,
-                self._sklearn_object,
-                ["snowflake-snowpark-python"] + self._get_dependencies(),
-                self.input_cols,
-                self.label_cols,
-                self.sample_weight_col,
-            )
+        model_trainer = ModelTrainerBuilder.build(
+            estimator=self._sklearn_object,
+            dataset=dataset,
+            input_cols=self.input_cols,
+            label_cols=self.label_cols,
+            sample_weight_col=self.sample_weight_col,
+            autogenerated=False,
+            subproject=_SUBPROJECT,
+        )
+        self._sklearn_object = model_trainer.train()
+        self._is_fitted = True
+        self._get_model_signatures(dataset)
+        return self
 
     def _get_pass_through_columns(self, dataset: DataFrame) -> List[str]:
         if self._drop_input_cols:
@@ -523,10 +488,6 @@ class GridSearchCV(BaseTransformer):
         project=_PROJECT,
         subproject=_SUBPROJECT,
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-    )
     def predict(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]:
         """Call predict on the estimator with the best found parameters
         For more details on this function, see [sklearn.model_selection.GridSearchCV.predict]
@@ -569,10 +530,6 @@ class GridSearchCV(BaseTransformer):
         project=_PROJECT,
         subproject=_SUBPROJECT,
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-    )
     def transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]:
         """Call transform on the estimator with the best found parameters
         For more details on this function, see [sklearn.model_selection.GridSearchCV.transform]
@@ -636,10 +593,6 @@ class GridSearchCV(BaseTransformer):
         project=_PROJECT,
         subproject=_SUBPROJECT,
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-    )
     def predict_proba(
         self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "predict_proba_"
     ) -> Union[DataFrame, pd.DataFrame]:
@@ -677,10 +630,6 @@ class GridSearchCV(BaseTransformer):
         project=_PROJECT,
         subproject=_SUBPROJECT,
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-    )
     def predict_log_proba(
         self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "predict_log_proba_"
     ) -> Union[DataFrame, pd.DataFrame]:
@@ -719,10 +668,6 @@ class GridSearchCV(BaseTransformer):
         project=_PROJECT,
         subproject=_SUBPROJECT,
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-    )
     def decision_function(
         self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "decision_function_"
     ) -> Union[DataFrame, pd.DataFrame]:
@@ -759,6 +704,8 @@ class GridSearchCV(BaseTransformer):
     @available_if(original_estimator_has_callable("score"))  # type: ignore[misc]
     def score(self, dataset: Union[DataFrame, pd.DataFrame]) -> float:
         """
+        If implemented by the original estimator, return the score for the dataset.
+
         Args:
             dataset: Union[snowflake.snowpark.DataFrame, pandas.DataFrame]
                 Snowpark or Pandas DataFrame.
@@ -811,9 +758,9 @@ 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 = _infer_signature(dataset[self.label_cols], "output")
+                outputs = list(_infer_signature(dataset[self.label_cols], "output"))
                 # rename the output columns
-                outputs = model_signature_utils.rename_features(outputs, self.output_cols)
+                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
                 )
@@ -850,6 +797,9 @@ class GridSearchCV(BaseTransformer):
         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