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

snowflake/ml/modeling/_internal/snowpark_trainer.py

@@ -0,0 +1,331 @@
+import importlib
+import inspect
+import os
+import posixpath
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import cloudpickle as cp
+
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.exceptions import (
+    error_codes,
+    exceptions,
+    modeling_error_messages,
+)
+from snowflake.ml._internal.utils import identifier, snowpark_dataframe_utils
+from snowflake.ml._internal.utils.query_result_checker import SqlResultValidator
+from snowflake.ml._internal.utils.temp_file_utils import (
+    cleanup_temp_files,
+    get_temp_file_path,
+)
+from snowflake.ml.modeling._internal.model_specifications import (
+    ModelSpecifications,
+    ModelSpecificationsBuilder,
+)
+from snowflake.snowpark import DataFrame, Session, exceptions as snowpark_exceptions
+from snowflake.snowpark._internal.utils import (
+    TempObjectType,
+    random_name_for_temp_object,
+)
+from snowflake.snowpark.functions import sproc
+from snowflake.snowpark.stored_procedure import StoredProcedure
+
+cp.register_pickle_by_value(inspect.getmodule(get_temp_file_path))
+cp.register_pickle_by_value(inspect.getmodule(identifier.get_inferred_name))
+
+_PROJECT = "ModelDevelopment"
+
+
+class SnowparkModelTrainer:
+    """
+    A class for training models on Snowflake data using the Sproc.
+
+    TODO (snandamuri): Introduce the concept of executor that would take the training function
+    and execute it on the target environments like, local, Snowflake warehouse, or SPCS, etc.
+    """
+
+    def __init__(
+        self,
+        estimator: object,
+        dataset: DataFrame,
+        session: Session,
+        input_cols: List[str],
+        label_cols: Optional[List[str]],
+        sample_weight_col: Optional[str],
+        autogenerated: bool = False,
+        subproject: str = "",
+    ) -> None:
+        """
+        Initializes the SnowparkModelTrainer with a model, a Snowpark DataFrame, feature, and label column names.
+
+        Args:
+            estimator: SKLearn compatible estimator or transformer object.
+            dataset: The dataset used for training the model.
+            session: Snowflake session object to be used for training.
+            input_cols: The name(s) of one or more columns in a DataFrame containing a feature to be used for training.
+            label_cols: The name(s) of one or more columns in a DataFrame representing the target variable(s) to learn.
+            sample_weight_col: The column name representing the weight of training examples.
+            autogenerated: A boolean denoting if the trainer is being used by autogenerated code or not.
+            subproject: subproject name to be used in telemetry.
+        """
+        self.estimator = estimator
+        self.dataset = dataset
+        self.session = session
+        self.input_cols = input_cols
+        self.label_cols = label_cols
+        self.sample_weight_col = sample_weight_col
+        self._autogenerated = autogenerated
+        self._subproject = subproject
+        self._class_name = estimator.__class__.__name__
+
+    def _create_temp_stage(self) -> str:
+        """
+        Creates temporary stage.
+
+        Returns:
+            Temp stage name.
+        """
+        # Create temp stage to upload pickled model file.
+        transform_stage_name = random_name_for_temp_object(TempObjectType.STAGE)
+        stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {transform_stage_name};"
+        SqlResultValidator(session=self.session, query=stage_creation_query).has_dimensions(
+            expected_rows=1, expected_cols=1
+        ).validate()
+        return transform_stage_name
+
+    def _upload_model_to_stage(self, stage_name: str) -> Tuple[str, str]:
+        """
+        Util method to pickle and upload the model to a temp Snowflake stage.
+
+        Args:
+            stage_name: Stage name to save model.
+
+        Returns:
+            a tuple containing stage file paths for pickled input model for training and location to store trained
+            models(response from training sproc).
+        """
+        # Create a temp file and dump the transform to that file.
+        local_transform_file_name = get_temp_file_path()
+        with open(local_transform_file_name, mode="w+b") as local_transform_file:
+            cp.dump(self.estimator, local_transform_file)
+
+        # Use posixpath to construct stage paths
+        stage_transform_file_name = posixpath.join(stage_name, os.path.basename(local_transform_file_name))
+        stage_result_file_name = posixpath.join(stage_name, os.path.basename(local_transform_file_name))
+
+        statement_params = telemetry.get_function_usage_statement_params(
+            project=_PROJECT,
+            subproject=self._subproject,
+            function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
+            api_calls=[sproc],
+            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+        )
+        # Put locally serialized transform on stage.
+        self.session.file.put(
+            local_transform_file_name,
+            stage_transform_file_name,
+            auto_compress=False,
+            overwrite=True,
+            statement_params=statement_params,
+        )
+
+        cleanup_temp_files([local_transform_file_name])
+        return (stage_transform_file_name, stage_result_file_name)
+
+    def _fetch_model_from_stage(self, dir_path: str, file_name: str, statement_params: Dict[str, str]) -> object:
+        """
+        Downloads the serialized model from a stage location and unpickels it.
+
+        Args:
+            dir_path: Stage directory path where results are stored.
+            file_name: File name with in the directory where results are stored.
+            statement_params: Statement params to be attached to the SQL queries issue form this method.
+
+        Returns:
+            Deserialized model object.
+        """
+        local_result_file_name = get_temp_file_path()
+        self.session.file.get(
+            posixpath.join(dir_path, file_name),
+            local_result_file_name,
+            statement_params=statement_params,
+        )
+
+        with open(os.path.join(local_result_file_name, file_name), mode="r+b") as result_file_obj:
+            fit_estimator = cp.load(result_file_obj)
+
+        cleanup_temp_files([local_result_file_name])
+        return fit_estimator
+
+    def _build_fit_wrapper_sproc(
+        self,
+        model_spec: ModelSpecifications,
+    ) -> Callable[[Any, List[str], str, str, List[str], List[str], Optional[str], Dict[str, str]], str]:
+        """
+        Constructs and returns a python stored procedure function to be used for training model.
+
+        Args:
+            model_spec: ModelSpecifications object that contains model specific information
+                like required imports, package dependencies, etc.
+
+        Returns:
+            A callable that can be registered as a stored procedure.
+        """
+        imports = model_spec.imports  # In order for the sproc to not resolve this reference in snowflake.ml
+
+        def fit_wrapper_function(
+            session: Session,
+            sql_queries: List[str],
+            stage_transform_file_name: str,
+            stage_result_file_name: str,
+            input_cols: List[str],
+            label_cols: List[str],
+            sample_weight_col: Optional[str],
+            statement_params: Dict[str, str],
+        ) -> str:
+            import inspect
+            import os
+
+            import cloudpickle as cp
+            import pandas as pd
+
+            for import_name in imports:
+                importlib.import_module(import_name)
+
+            # Execute snowpark queries and obtain the results as pandas dataframe
+            # NB: this implies that the result data must fit into memory.
+            for query in sql_queries[:-1]:
+                _ = session.sql(query).collect(statement_params=statement_params)
+            sp_df = session.sql(sql_queries[-1])
+            df: pd.DataFrame = sp_df.to_pandas(statement_params=statement_params)
+            df.columns = sp_df.columns
+
+            local_transform_file_name = get_temp_file_path()
+
+            session.file.get(stage_transform_file_name, local_transform_file_name, statement_params=statement_params)
+
+            local_transform_file_path = os.path.join(
+                local_transform_file_name, os.listdir(local_transform_file_name)[0]
+            )
+            with open(local_transform_file_path, mode="r+b") as local_transform_file_obj:
+                estimator = cp.load(local_transform_file_obj)
+
+            argspec = inspect.getfullargspec(estimator.fit)
+            args = {"X": df[input_cols]}
+            if label_cols:
+                label_arg_name = "Y" if "Y" in argspec.args else "y"
+                args[label_arg_name] = df[label_cols].squeeze()
+
+            if sample_weight_col is not None and "sample_weight" in argspec.args:
+                args["sample_weight"] = df[sample_weight_col].squeeze()
+
+            estimator.fit(**args)
+
+            local_result_file_name = get_temp_file_path()
+
+            with open(local_result_file_name, mode="w+b") as local_result_file_obj:
+                cp.dump(estimator, local_result_file_obj)
+
+            session.file.put(
+                local_result_file_name,
+                stage_result_file_name,
+                auto_compress=False,
+                overwrite=True,
+                statement_params=statement_params,
+            )
+
+            # Note: you can add something like  + "|" + str(df) to the return string
+            # to pass debug information to the caller.
+            return str(os.path.basename(local_result_file_name))
+
+        return fit_wrapper_function
+
+    def _get_fit_wrapper_sproc(self, statement_params: Dict[str, str]) -> StoredProcedure:
+        # If the sproc already exists, don't register.
+        if not hasattr(self.session, "_FIT_WRAPPER_SPROCS"):
+            self.session._FIT_WRAPPER_SPROCS: Dict[str, StoredProcedure] = {}  # type: ignore[attr-defined, misc]
+
+        model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
+        fit_sproc_key = model_spec.__class__.__name__
+        if fit_sproc_key in self.session._FIT_WRAPPER_SPROCS:  # type: ignore[attr-defined]
+            fit_sproc: StoredProcedure = self.session._FIT_WRAPPER_SPROCS[fit_sproc_key]  # type: ignore[attr-defined]
+            return fit_sproc
+
+        fit_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
+
+        fit_wrapper_sproc = self.session.sproc.register(
+            func=self._build_fit_wrapper_sproc(model_spec=model_spec),
+            is_permanent=False,
+            name=fit_sproc_name,
+            packages=["snowflake-snowpark-python"] + model_spec.pkgDependencies,  # type: ignore[arg-type]
+            replace=True,
+            session=self.session,
+            statement_params=statement_params,
+        )
+
+        self.session._FIT_WRAPPER_SPROCS[fit_sproc_key] = fit_wrapper_sproc  # type: ignore[attr-defined]
+
+        return fit_wrapper_sproc
+
+    def train(self) -> object:
+        """
+        Trains the model by pushing down the compute into Snowflake using stored procedures.
+
+        Returns:
+            Trained model
+
+        Raises:
+            e: Raises an exception if any of Snowflake operations fail because of any reason.
+            SnowflakeMLException: Know exception are caught and rethrow with more detailed error message.
+        """
+        dataset = snowpark_dataframe_utils.cast_snowpark_dataframe_column_types(self.dataset)
+
+        # TODO(snandamuri) : Handle the already in a stored procedure case in the in builder.
+
+        # Extract query that generated the dataframe. We will need to pass it to the fit procedure.
+        queries = dataset.queries["queries"]
+
+        transform_stage_name = self._create_temp_stage()
+        (stage_transform_file_name, stage_result_file_name) = self._upload_model_to_stage(
+            stage_name=transform_stage_name
+        )
+
+        # Call fit sproc
+        statement_params = telemetry.get_function_usage_statement_params(
+            project=_PROJECT,
+            subproject=self._subproject,
+            function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
+            api_calls=[Session.call],
+            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+        )
+
+        fit_wrapper_sproc = self._get_fit_wrapper_sproc(statement_params=statement_params)
+
+        try:
+            sproc_export_file_name: str = fit_wrapper_sproc(
+                self.session,
+                queries,
+                stage_transform_file_name,
+                stage_result_file_name,
+                self.input_cols,
+                self.label_cols,
+                self.sample_weight_col,
+                statement_params,
+            )
+        except snowpark_exceptions.SnowparkClientException as e:
+            if "fit() missing 1 required positional argument: 'y'" in str(e):
+                raise exceptions.SnowflakeMLException(
+                    error_code=error_codes.NOT_FOUND,
+                    original_exception=RuntimeError(modeling_error_messages.ATTRIBUTE_NOT_SET.format("label_cols")),
+                ) from e
+            raise e
+
+        if "|" in sproc_export_file_name:
+            fields = sproc_export_file_name.strip().split("|")
+            sproc_export_file_name = fields[0]
+
+        return self._fetch_model_from_stage(
+            dir_path=stage_result_file_name,
+            file_name=sproc_export_file_name,
+            statement_params=statement_params,
+        )

snowflake/ml/modeling/calibration/calibrated_classifier_cv.py

@@ -22,17 +22,19 @@ from sklearn.utils.metaestimators import available_if
 from snowflake.ml.modeling.framework.base import BaseTransformer, _process_cols
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions, modeling_error_messages
+from snowflake.ml._internal.env_utils import SNOWML_SPROC_ENV
 from snowflake.ml._internal.utils import pkg_version_utils, identifier
-from snowflake.snowpark import DataFrame
+from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.ml.modeling._internal.snowpark_handlers import SnowparkHandlers as HandlersImpl
+from snowflake.ml.modeling._internal.model_trainer_builder import ModelTrainerBuilder
+from snowflake.ml.modeling._internal.model_trainer import ModelTrainer
 from snowflake.ml.modeling._internal.estimator_utils import (
     gather_dependencies,
     original_estimator_has_callable,
     transform_snowml_obj_to_sklearn_obj,
     validate_sklearn_args,
 )
-from snowflake.ml.modeling._internal.snowpark_handlers import SklearnWrapperProvider
 from snowflake.ml.modeling._internal.estimator_protocols import FitPredictHandlers
 
 from snowflake.ml.model.model_signature import (
@@ -52,7 +54,6 @@ _PROJECT = "ModelDevelopment"
 _SUBPROJECT = "".join([s.capitalize() for s in "sklearn.calibration".replace("sklearn.", "").split("_")])
 
 
-
 class CalibratedClassifierCV(BaseTransformer):
     r"""Probability calibration with isotonic regression or logistic regression
     For more details on this class, see [sklearn.calibration.CalibratedClassifierCV]
@@ -60,6 +61,51 @@ class CalibratedClassifierCV(BaseTransformer):
 
     Parameters
     ----------
+
+    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, sample_weight_col, and passthrough_cols
+        parameters are considered input columns. Input columns can also be set after
+        initialization with the `set_input_cols` method.
+    
+    label_cols: Optional[Union[str, List[str]]]
+        A string or list of strings representing column names that contain labels.
+        Label columns must be specified with this parameter during initialization
+        or with the `set_label_cols` method before fitting.
+
+    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 must
+        match the expected number of output columns from the specific predictor or
+        transformer class used.
+        If you omit this parameter, output column names are derived by adding an
+        OUTPUT_ prefix to the label column names for supervised estimators, or
+        OUTPUT_<IDX>for unsupervised estimators. These inferred output column names
+        work for predictors, but output_cols must be set explicitly for transformers.
+        In general, explicitly specifying output column names is clearer, especially
+        if you don’t specify the input column names.
+        To transform in place, pass the same names for input_cols and output_cols.
+        be set explicitly for transformers. Output columns can also be set after
+        initialization with the `set_output_cols` method.
+
+    sample_weight_col: Optional[str]
+        A string representing the column name containing the sample weights.
+        This argument is only required when working with weighted datasets. Sample
+        weight column can also be set after initialization with the
+        `set_sample_weight_col` method.
+
+    passthrough_cols: Optional[Union[str, List[str]]]
+        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. Passthrough columns
+        can also be set after initialization with the `set_passthrough_cols` method.
+
+    drop_input_cols: Optional[bool], default=False
+        If set, the response of predict(), transform() methods will not contain input columns.
+
     estimator: estimator instance, default=None
         The classifier whose output need to be calibrated to provide more
         accurate `predict_proba` outputs. The default classifier is
@@ -121,42 +167,6 @@ class CalibratedClassifierCV(BaseTransformer):
 
     base_estimator: estimator instance
         This parameter is deprecated. Use `estimator` instead.
-
-    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, sample_weight_col, and passthrough_cols
-        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 must
-        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.
-
-    sample_weight_col: Optional[str]
-        A string representing the column name containing the sample weights.
-        This argument is only required when working with weighted datasets.
-
-    passthrough_cols: Optional[Union[str, List[str]]]
-        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.
-
-    drop_input_cols: Optional[bool], default=False
-        If set, the response of predict(), transform() methods will not contain input columns.
     """
 
     def __init__(  # type: ignore[no-untyped-def]
@@ -183,7 +193,7 @@ class CalibratedClassifierCV(BaseTransformer):
         self.set_passthrough_cols(passthrough_cols)
         self.set_drop_input_cols(drop_input_cols)
         self.set_sample_weight_col(sample_weight_col)        
-        deps = set(SklearnWrapperProvider().dependencies)
+        deps: Set[str] = set([f'numpy=={np.__version__}', f'scikit-learn=={sklearn.__version__}', f'cloudpickle=={cp.__version__}'])
         deps = deps | gather_dependencies(estimator)
         deps = deps | gather_dependencies(base_estimator)
         self._deps = list(deps)
@@ -199,13 +209,14 @@ class CalibratedClassifierCV(BaseTransformer):
             args=init_args,
             klass=sklearn.calibration.CalibratedClassifierCV
         )
-        self._sklearn_object = sklearn.calibration.CalibratedClassifierCV(
+        self._sklearn_object: Any = sklearn.calibration.CalibratedClassifierCV(
             **cleaned_up_init_args,
         )
         self._model_signature_dict: Optional[Dict[str, ModelSignature]] = None
         # If user used snowpark dataframe during fit, here it stores the snowpark input_cols, otherwise the processed input_cols
         self._snowpark_cols: Optional[List[str]] = self.input_cols
-        self._handlers: FitPredictHandlers = HandlersImpl(class_name=CalibratedClassifierCV.__class__.__name__, subproject=_SUBPROJECT, autogenerated=True, wrapper_provider=SklearnWrapperProvider())
+        self._handlers: FitPredictHandlers = HandlersImpl(class_name=CalibratedClassifierCV.__class__.__name__, subproject=_SUBPROJECT, autogenerated=True)
+        self._autogenerated = True
 
     def _get_rand_id(self) -> str:
         """
@@ -261,54 +272,48 @@ class CalibratedClassifierCV(BaseTransformer):
             self
         """
         self._infer_input_output_cols(dataset)
-        if isinstance(dataset, pd.DataFrame):
-            assert self._sklearn_object is not None  # keep mypy happy
-            self._sklearn_object = 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 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)
+            
+            # 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)
+
+            self._snowpark_cols = dataset.select(self.input_cols).columns
+
+             # If we are already in a stored procedure, no need to kick off another one.
+            if SNOWML_SPROC_ENV in os.environ:
+                statement_params = telemetry.get_function_usage_statement_params(
+                    project=_PROJECT,
+                    subproject=_SUBPROJECT,
+                    function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), CalibratedClassifierCV.__class__.__name__),
+                    api_calls=[Session.call],
+                    custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
+                )
+                pd_df: pd.DataFrame = dataset.to_pandas(statement_params=statement_params)
+                pd_df.columns = dataset.columns
+                dataset = pd_df
+
+        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=self._autogenerated,
+            subproject=_SUBPROJECT
+        )
+        self._sklearn_object = model_trainer.train()
         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)
-
-        estimator = self._sklearn_object
-        assert estimator is not None  # Keep mypy happy
-
-        self._snowpark_cols = dataset.select(self.input_cols).columns
-
-        self._sklearn_object = self._handlers.fit_snowpark(
-            dataset,
-            session,
-            estimator,
-            ["snowflake-snowpark-python"] + self._get_dependencies(),
-            self.input_cols,
-            self.label_cols,
-            self.sample_weight_col,
-        )
-
     def _get_pass_through_columns(self, dataset: DataFrame) -> List[str]:
         if self._drop_input_cols:
             return []
@@ -496,11 +501,6 @@ class CalibratedClassifierCV(BaseTransformer):
         subproject=_SUBPROJECT,
         custom_tags=dict([("autogen", True)]),
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
     def predict(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]:
         """Predict the target of new samples
         For more details on this function, see [sklearn.calibration.CalibratedClassifierCV.predict]
@@ -554,11 +554,6 @@ class CalibratedClassifierCV(BaseTransformer):
         subproject=_SUBPROJECT,
         custom_tags=dict([("autogen", True)]),
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
     def transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]:
         """Method not supported for this class.
 
@@ -615,7 +610,8 @@ class CalibratedClassifierCV(BaseTransformer):
         if False:
             self.fit(dataset)
             assert self._sklearn_object is not None
-            return self._sklearn_object.labels_
+            labels : npt.NDArray[Any] = self._sklearn_object.labels_
+            return labels
         else:
             raise NotImplementedError
 
@@ -651,6 +647,7 @@ class CalibratedClassifierCV(BaseTransformer):
             output_cols = []
 
         # Make sure column names are valid snowflake identifiers.
+        assert output_cols is not None  # Make MyPy happy
         rv = [identifier.rename_to_valid_snowflake_identifier(c) for c in output_cols]
 
         return rv
@@ -661,11 +658,6 @@ class CalibratedClassifierCV(BaseTransformer):
         subproject=_SUBPROJECT,
         custom_tags=dict([("autogen", True)]),
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
     def predict_proba(
         self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "predict_proba_"
     ) -> Union[DataFrame, pd.DataFrame]:
@@ -708,11 +700,6 @@ class CalibratedClassifierCV(BaseTransformer):
         subproject=_SUBPROJECT,
         custom_tags=dict([("autogen", True)]),
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
     def predict_log_proba(
         self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "predict_log_proba_"
     ) -> Union[DataFrame, pd.DataFrame]:
@@ -751,16 +738,6 @@ class CalibratedClassifierCV(BaseTransformer):
         return output_df
 
     @available_if(original_estimator_has_callable("decision_function"))  # type: ignore[misc]
-    @telemetry.send_api_usage_telemetry(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
     def decision_function(
         self, dataset: Union[DataFrame, pd.DataFrame], output_cols_prefix: str = "decision_function_"
     ) -> Union[DataFrame, pd.DataFrame]:
@@ -861,11 +838,6 @@ class CalibratedClassifierCV(BaseTransformer):
         subproject=_SUBPROJECT,
         custom_tags=dict([("autogen", True)]),
     )
-    @telemetry.add_stmt_params_to_df(
-        project=_PROJECT,
-        subproject=_SUBPROJECT,
-        custom_tags=dict([("autogen", True)]),
-    )
     def kneighbors(
         self,
         dataset: Union[DataFrame, pd.DataFrame],
@@ -925,9 +897,9 @@ class CalibratedClassifierCV(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", use_snowflake_identifiers=True)
+                outputs = list(_infer_signature(dataset[self.label_cols], "output", use_snowflake_identifiers=True))
                 # 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)