snowflake-ml-python 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

snowflake/ml/model/type_hints.py

@@ -9,6 +9,8 @@ if TYPE_CHECKING:
     import pandas as pd
     import sklearn.base
     import sklearn.pipeline
+    import tensorflow
+    import torch
     import xgboost
 
     import snowflake.ml.model.custom_model
@@ -34,9 +36,10 @@ _SupportedNumpyDtype = Union[
 ]
 _SupportedNumpyArray = npt.NDArray[_SupportedNumpyDtype]
 _SupportedBuiltinsList = Sequence[_SupportedBuiltins]
+_SupportedArrayLike = Union[_SupportedNumpyArray, "torch.Tensor", "tensorflow.Tensor", "tensorflow.Variable"]
 
 SupportedLocalDataType = Union[
-    "pd.DataFrame", _SupportedNumpyArray, Sequence[_SupportedNumpyArray], _SupportedBuiltinsList
+    "pd.DataFrame", _SupportedNumpyArray, Sequence[_SupportedArrayLike], _SupportedBuiltinsList
 ]
 
 SupportedDataType = Union[SupportedLocalDataType, "snowflake.snowpark.DataFrame"]
@@ -51,6 +54,8 @@ SupportedLocalModelType = Union[
     "sklearn.pipeline.Pipeline",
     "xgboost.XGBModel",
     "xgboost.Booster",
+    "torch.nn.Module",
+    "torch.jit.ScriptModule",  # type:ignore[name-defined]
 ]
 
 SupportedSnowMLModelType: TypeAlias = "base.BaseEstimator"
@@ -70,6 +75,8 @@ Here is all acceptable types of Snowflake native model packaging and its handler
 | xgboost.XGBModel       | xgboost.py   | _XGBModelHandler    |
 | xgboost.Booster        | xgboost.py   | _XGBModelHandler    |
 | snowflake.ml.framework.base.BaseEstimator      | snowmlmodel.py   | _SnowMLModelHandler    |
+| torch.nn.Module      | pytroch.py   | _PyTorchHandler    |
+| torch.jit.ScriptModule      | torchscript.py   | _TorchScripthHandler    |
 """
 
 
@@ -79,19 +86,23 @@ _ModelType = TypeVar("_ModelType", bound=SupportedModelType)
 class DeployOptions(TypedDict):
     """Common Options for deploying to Snowflake.
 
-    output_with_input_features: Whether or not preserve the input columns in the output when predicting.
-        Defaults to False.
+    disable_local_conda_resolver: Set to disable use local conda resolver to do pre-check on environment and rely on
+        the information schema only. Defaults to False.
     keep_order: Whether or not preserve the row order when predicting. Only available for dataframe has fewer than 2**64
         rows. Defaults to True.
+    output_with_input_features: Whether or not preserve the input columns in the output when predicting.
+        Defaults to False.
     """
 
-    output_with_input_features: NotRequired[bool]
+    disable_local_conda_resolver: NotRequired[bool]
     keep_order: NotRequired[bool]
+    output_with_input_features: NotRequired[bool]
 
 
 class WarehouseDeployOptions(DeployOptions):
     """Options for deploying to the Snowflake Warehouse.
 
+
     permanent_udf_stage_location: A Snowflake stage option where the UDF should be persisted. If specified, the model
         will be deployed as a permanent UDF, otherwise temporary.
     relax_version: Whether or not relax the version constraints of the dependencies if unresolvable. Defaults to False.
@@ -130,3 +141,11 @@ class XGBModelSaveOptions(ModelSaveOption):
 
 class SNOWModelSaveOptions(ModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
+
+
+class PyTorchSaveOptions(ModelSaveOption):
+    target_methods: NotRequired[Sequence[str]]
+
+
+class TorchScriptSaveOptions(ModelSaveOption):
+    target_methods: NotRequired[Sequence[str]]

snowflake/ml/modeling/calibration/calibrated_classifier_cv.py

@@ -7,6 +7,7 @@
 #
 import inspect
 import os
+import posixpath
 from typing import Iterable, Optional, Union, List, Any, Dict, Callable, Set
 from uuid import uuid4
 
@@ -27,6 +28,7 @@ from snowflake.ml._internal.utils.temp_file_utils import cleanup_temp_files, get
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark.functions import pandas_udf, sproc
 from snowflake.snowpark.types import PandasSeries
+from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 
 from snowflake.ml.model.model_signature import (
     DataType,
@@ -247,7 +249,6 @@ class CalibratedClassifierCV(BaseTransformer):
         sample_weight_col: Optional[str] = None,
     ) -> None:
         super().__init__()
-        self.id = str(uuid4()).replace("-", "_").upper()
         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)
@@ -274,6 +275,15 @@ class CalibratedClassifierCV(BaseTransformer):
         self.set_drop_input_cols(drop_input_cols)
         self.set_sample_weight_col(sample_weight_col)
 
+    def _get_rand_id(self) -> str:
+        """
+        Generate random id to be used in sproc and stage names.
+
+        Returns:
+            Random id string usable in sproc, table, and stage names.
+        """
+        return str(uuid4()).replace("-", "_").upper()
+
     def _infer_input_output_cols(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
         """
         Infer `self.input_cols` and `self.output_cols` if they are not explicitly set.
@@ -352,7 +362,7 @@ class CalibratedClassifierCV(BaseTransformer):
             cp.dump(self._sklearn_object, local_transform_file)
 
         # Create temp stage to run fit.
-        transform_stage_name = "SNOWML_TRANSFORM_{safe_id}".format(safe_id=self.id)
+        transform_stage_name = "SNOWML_TRANSFORM_{safe_id}".format(safe_id=self._get_rand_id())
         stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {transform_stage_name};"
         SqlResultValidator(
             session=session,
@@ -365,11 +375,12 @@ class CalibratedClassifierCV(BaseTransformer):
             expected_value=f"Stage area {transform_stage_name} successfully created."
         ).validate()
 
-        stage_transform_file_name = os.path.join(transform_stage_name, os.path.basename(local_transform_file_name))
+        # Use posixpath to construct stage paths
+        stage_transform_file_name = posixpath.join(transform_stage_name, os.path.basename(local_transform_file_name))
+        stage_result_file_name = posixpath.join(transform_stage_name, os.path.basename(local_transform_file_name))
         local_result_file_name = get_temp_file_path()
-        stage_result_file_name = os.path.join(transform_stage_name, os.path.basename(local_transform_file_name))
 
-        fit_sproc_name = "SNOWML_FIT_{safe_id}".format(safe_id=self.id)
+        fit_sproc_name = "SNOWML_FIT_{safe_id}".format(safe_id=self._get_rand_id())
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=_SUBPROJECT,
@@ -395,6 +406,7 @@ class CalibratedClassifierCV(BaseTransformer):
             replace=True,
             session=session,
             statement_params=statement_params,
+            anonymous=True
         )
         def fit_wrapper_sproc(
             session: Session,
@@ -403,7 +415,8 @@ class CalibratedClassifierCV(BaseTransformer):
             stage_result_file_name: str,
             input_cols: List[str],
             label_cols: List[str],
-            sample_weight_col: Optional[str]
+            sample_weight_col: Optional[str],
+            statement_params: Dict[str, str]
         ) -> str:
             import cloudpickle as cp
             import numpy as np
@@ -470,15 +483,15 @@ class CalibratedClassifierCV(BaseTransformer):
             api_calls=[Session.call],
             custom_tags=dict([("autogen", True)]),
         )
-        sproc_export_file_name = session.call(
-            fit_sproc_name,
+        sproc_export_file_name = fit_wrapper_sproc(
+            session,
             query,
             stage_transform_file_name,
             stage_result_file_name,
             identifier.get_unescaped_names(self.input_cols),
             identifier.get_unescaped_names(self.label_cols),
             identifier.get_unescaped_names(self.sample_weight_col),
-            statement_params=statement_params,
+            statement_params,
         )
 
         if "|" in sproc_export_file_name:
@@ -488,7 +501,7 @@ class CalibratedClassifierCV(BaseTransformer):
                 print("\n".join(fields[1:]))
 
         session.file.get(
-            os.path.join(stage_result_file_name, sproc_export_file_name),
+            posixpath.join(stage_result_file_name, sproc_export_file_name),
             local_result_file_name,
             statement_params=statement_params
         )
@@ -534,7 +547,7 @@ class CalibratedClassifierCV(BaseTransformer):
 
         # Register vectorized UDF for batch inference
         batch_inference_udf_name = "SNOWML_BATCH_INFERENCE_{safe_id}_{method}".format(
-                safe_id=self.id, method=inference_method)
+                safe_id=self._get_rand_id(), method=inference_method)
 
         # Need to do this since if we use self._sklearn_object directly in the UDF, Snowpark
         # will try to pickle all of self which fails.
@@ -626,7 +639,7 @@ class CalibratedClassifierCV(BaseTransformer):
             return transformed_pandas_df.to_dict("records")
 
         batch_inference_table_name = "SNOWML_BATCH_INFERENCE_INPUT_TABLE_{safe_id}".format(
-            safe_id=self.id
+            safe_id=self._get_rand_id()
         )
 
         pass_through_columns = self._get_pass_through_columns(dataset)
@@ -793,11 +806,18 @@ class CalibratedClassifierCV(BaseTransformer):
             Transformed dataset.
         """
         if isinstance(dataset, DataFrame):
+            expected_type_inferred = ""
+            # when it is classifier, infer the datatype from label columns
+            if expected_type_inferred == "" and 'predict' in self.model_signatures:
+                expected_type_inferred = convert_sp_to_sf_type(
+                    self.model_signatures['predict'].outputs[0].as_snowpark_type()
+                )
+            
             output_df = self._batch_inference(
                 dataset=dataset,
                 inference_method="predict",
                 expected_output_cols_list=self.output_cols,
-                expected_output_cols_type="",
+                expected_output_cols_type=expected_type_inferred,
             )
         elif isinstance(dataset, pd.DataFrame):
             output_df = self._sklearn_inference(
@@ -868,10 +888,10 @@ class CalibratedClassifierCV(BaseTransformer):
 
     def _get_output_column_names(self, output_cols_prefix: str) -> List[str]:
         """ Returns the list of output columns for predict_proba(), decision_function(), etc.. functions.
-        Returns an empty list if current object is not a classifier or not yet fitted.
+        Returns a list with output_cols_prefix as the only element if the estimator is not a classifier.
         """
         if getattr(self._sklearn_object, "classes_", None) is None:
-            return []
+            return [output_cols_prefix]
 
         classes = self._sklearn_object.classes_
         if isinstance(classes, numpy.ndarray):
@@ -1100,7 +1120,7 @@ class CalibratedClassifierCV(BaseTransformer):
             cp.dump(self._sklearn_object, local_score_file)
 
         # Create temp stage to run score.
-        score_stage_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self.id)
+        score_stage_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self._get_rand_id())
         session = dataset._session
         stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {score_stage_name};"
         SqlResultValidator(
@@ -1114,8 +1134,9 @@ class CalibratedClassifierCV(BaseTransformer):
             expected_value=f"Stage area {score_stage_name} successfully created."
         ).validate()
 
-        stage_score_file_name = os.path.join(score_stage_name, os.path.basename(local_score_file_name))
-        score_sproc_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self.id)
+        # Use posixpath to construct stage paths
+        stage_score_file_name = posixpath.join(score_stage_name, os.path.basename(local_score_file_name))
+        score_sproc_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self._get_rand_id())
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=_SUBPROJECT,
@@ -1141,6 +1162,7 @@ class CalibratedClassifierCV(BaseTransformer):
             replace=True,
             session=session,
             statement_params=statement_params,
+            anonymous=True
         )
         def score_wrapper_sproc(
             session: Session,
@@ -1148,7 +1170,8 @@ class CalibratedClassifierCV(BaseTransformer):
             stage_score_file_name: str,
             input_cols: List[str],
             label_cols: List[str],
-            sample_weight_col: Optional[str]
+            sample_weight_col: Optional[str],
+            statement_params: Dict[str, str]
         ) -> float:
             import cloudpickle as cp
             import numpy as np
@@ -1198,14 +1221,14 @@ class CalibratedClassifierCV(BaseTransformer):
             api_calls=[Session.call],
             custom_tags=dict([("autogen", True)]),
         )
-        score = session.call(
-            score_sproc_name,
+        score = score_wrapper_sproc(
+            session,
             query,
             stage_score_file_name,
             identifier.get_unescaped_names(self.input_cols),
             identifier.get_unescaped_names(self.label_cols),
             identifier.get_unescaped_names(self.sample_weight_col),
-            statement_params=statement_params,
+            statement_params,
         )
 
         cleanup_temp_files([local_score_file_name])
@@ -1223,18 +1246,20 @@ class CalibratedClassifierCV(BaseTransformer):
             if self._sklearn_object._estimator_type == 'classifier':
                 outputs = _infer_signature(dataset[self.label_cols], "output")  # label columns is the desired type for output
                 outputs = _rename_features(outputs, self.output_cols)  # rename the output columns
-                self._model_signature_dict["predict"] = ModelSignature(inputs, outputs)
+                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, outputs)
-
+                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}_"
                 output_column_names = self._get_output_column_names(output_cols_prefix)
                 outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in output_column_names]
-                self._model_signature_dict[prob_func] = ModelSignature(inputs, outputs)
+                self._model_signature_dict[prob_func] = ModelSignature(inputs, 
+                                                                       ([] if self._drop_input_cols else inputs) + outputs)
 
     @property
     def model_signatures(self) -> Dict[str, ModelSignature]:

snowflake/ml/modeling/cluster/affinity_propagation.py

@@ -7,6 +7,7 @@
 #
 import inspect
 import os
+import posixpath
 from typing import Iterable, Optional, Union, List, Any, Dict, Callable, Set
 from uuid import uuid4
 
@@ -27,6 +28,7 @@ from snowflake.ml._internal.utils.temp_file_utils import cleanup_temp_files, get
 from snowflake.snowpark import DataFrame, Session
 from snowflake.snowpark.functions import pandas_udf, sproc
 from snowflake.snowpark.types import PandasSeries
+from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 
 from snowflake.ml.model.model_signature import (
     DataType,
@@ -224,7 +226,6 @@ class AffinityPropagation(BaseTransformer):
         sample_weight_col: Optional[str] = None,
     ) -> None:
         super().__init__()
-        self.id = str(uuid4()).replace("-", "_").upper()
         deps: Set[str] = set([f'numpy=={np.__version__}', f'scikit-learn=={sklearn.__version__}', f'cloudpickle=={cp.__version__}'])
         
         self._deps = list(deps)
@@ -251,6 +252,15 @@ class AffinityPropagation(BaseTransformer):
         self.set_drop_input_cols(drop_input_cols)
         self.set_sample_weight_col(sample_weight_col)
 
+    def _get_rand_id(self) -> str:
+        """
+        Generate random id to be used in sproc and stage names.
+
+        Returns:
+            Random id string usable in sproc, table, and stage names.
+        """
+        return str(uuid4()).replace("-", "_").upper()
+
     def _infer_input_output_cols(self, dataset: Union[DataFrame, pd.DataFrame]) -> None:
         """
         Infer `self.input_cols` and `self.output_cols` if they are not explicitly set.
@@ -329,7 +339,7 @@ class AffinityPropagation(BaseTransformer):
             cp.dump(self._sklearn_object, local_transform_file)
 
         # Create temp stage to run fit.
-        transform_stage_name = "SNOWML_TRANSFORM_{safe_id}".format(safe_id=self.id)
+        transform_stage_name = "SNOWML_TRANSFORM_{safe_id}".format(safe_id=self._get_rand_id())
         stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {transform_stage_name};"
         SqlResultValidator(
             session=session,
@@ -342,11 +352,12 @@ class AffinityPropagation(BaseTransformer):
             expected_value=f"Stage area {transform_stage_name} successfully created."
         ).validate()
 
-        stage_transform_file_name = os.path.join(transform_stage_name, os.path.basename(local_transform_file_name))
+        # Use posixpath to construct stage paths
+        stage_transform_file_name = posixpath.join(transform_stage_name, os.path.basename(local_transform_file_name))
+        stage_result_file_name = posixpath.join(transform_stage_name, os.path.basename(local_transform_file_name))
         local_result_file_name = get_temp_file_path()
-        stage_result_file_name = os.path.join(transform_stage_name, os.path.basename(local_transform_file_name))
 
-        fit_sproc_name = "SNOWML_FIT_{safe_id}".format(safe_id=self.id)
+        fit_sproc_name = "SNOWML_FIT_{safe_id}".format(safe_id=self._get_rand_id())
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=_SUBPROJECT,
@@ -372,6 +383,7 @@ class AffinityPropagation(BaseTransformer):
             replace=True,
             session=session,
             statement_params=statement_params,
+            anonymous=True
         )
         def fit_wrapper_sproc(
             session: Session,
@@ -380,7 +392,8 @@ class AffinityPropagation(BaseTransformer):
             stage_result_file_name: str,
             input_cols: List[str],
             label_cols: List[str],
-            sample_weight_col: Optional[str]
+            sample_weight_col: Optional[str],
+            statement_params: Dict[str, str]
         ) -> str:
             import cloudpickle as cp
             import numpy as np
@@ -447,15 +460,15 @@ class AffinityPropagation(BaseTransformer):
             api_calls=[Session.call],
             custom_tags=dict([("autogen", True)]),
         )
-        sproc_export_file_name = session.call(
-            fit_sproc_name,
+        sproc_export_file_name = fit_wrapper_sproc(
+            session,
             query,
             stage_transform_file_name,
             stage_result_file_name,
             identifier.get_unescaped_names(self.input_cols),
             identifier.get_unescaped_names(self.label_cols),
             identifier.get_unescaped_names(self.sample_weight_col),
-            statement_params=statement_params,
+            statement_params,
         )
 
         if "|" in sproc_export_file_name:
@@ -465,7 +478,7 @@ class AffinityPropagation(BaseTransformer):
                 print("\n".join(fields[1:]))
 
         session.file.get(
-            os.path.join(stage_result_file_name, sproc_export_file_name),
+            posixpath.join(stage_result_file_name, sproc_export_file_name),
             local_result_file_name,
             statement_params=statement_params
         )
@@ -511,7 +524,7 @@ class AffinityPropagation(BaseTransformer):
 
         # Register vectorized UDF for batch inference
         batch_inference_udf_name = "SNOWML_BATCH_INFERENCE_{safe_id}_{method}".format(
-                safe_id=self.id, method=inference_method)
+                safe_id=self._get_rand_id(), method=inference_method)
 
         # Need to do this since if we use self._sklearn_object directly in the UDF, Snowpark
         # will try to pickle all of self which fails.
@@ -603,7 +616,7 @@ class AffinityPropagation(BaseTransformer):
             return transformed_pandas_df.to_dict("records")
 
         batch_inference_table_name = "SNOWML_BATCH_INFERENCE_INPUT_TABLE_{safe_id}".format(
-            safe_id=self.id
+            safe_id=self._get_rand_id()
         )
 
         pass_through_columns = self._get_pass_through_columns(dataset)
@@ -770,11 +783,18 @@ class AffinityPropagation(BaseTransformer):
             Transformed dataset.
         """
         if isinstance(dataset, DataFrame):
+            expected_type_inferred = ""
+            # when it is classifier, infer the datatype from label columns
+            if expected_type_inferred == "" and 'predict' in self.model_signatures:
+                expected_type_inferred = convert_sp_to_sf_type(
+                    self.model_signatures['predict'].outputs[0].as_snowpark_type()
+                )
+            
             output_df = self._batch_inference(
                 dataset=dataset,
                 inference_method="predict",
                 expected_output_cols_list=self.output_cols,
-                expected_output_cols_type="",
+                expected_output_cols_type=expected_type_inferred,
             )
         elif isinstance(dataset, pd.DataFrame):
             output_df = self._sklearn_inference(
@@ -845,10 +865,10 @@ class AffinityPropagation(BaseTransformer):
 
     def _get_output_column_names(self, output_cols_prefix: str) -> List[str]:
         """ Returns the list of output columns for predict_proba(), decision_function(), etc.. functions.
-        Returns an empty list if current object is not a classifier or not yet fitted.
+        Returns a list with output_cols_prefix as the only element if the estimator is not a classifier.
         """
         if getattr(self._sklearn_object, "classes_", None) is None:
-            return []
+            return [output_cols_prefix]
 
         classes = self._sklearn_object.classes_
         if isinstance(classes, numpy.ndarray):
@@ -1073,7 +1093,7 @@ class AffinityPropagation(BaseTransformer):
             cp.dump(self._sklearn_object, local_score_file)
 
         # Create temp stage to run score.
-        score_stage_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self.id)
+        score_stage_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self._get_rand_id())
         session = dataset._session
         stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {score_stage_name};"
         SqlResultValidator(
@@ -1087,8 +1107,9 @@ class AffinityPropagation(BaseTransformer):
             expected_value=f"Stage area {score_stage_name} successfully created."
         ).validate()
 
-        stage_score_file_name = os.path.join(score_stage_name, os.path.basename(local_score_file_name))
-        score_sproc_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self.id)
+        # Use posixpath to construct stage paths
+        stage_score_file_name = posixpath.join(score_stage_name, os.path.basename(local_score_file_name))
+        score_sproc_name = "SNOWML_SCORE_{safe_id}".format(safe_id=self._get_rand_id())
         statement_params = telemetry.get_function_usage_statement_params(
             project=_PROJECT,
             subproject=_SUBPROJECT,
@@ -1114,6 +1135,7 @@ class AffinityPropagation(BaseTransformer):
             replace=True,
             session=session,
             statement_params=statement_params,
+            anonymous=True
         )
         def score_wrapper_sproc(
             session: Session,
@@ -1121,7 +1143,8 @@ class AffinityPropagation(BaseTransformer):
             stage_score_file_name: str,
             input_cols: List[str],
             label_cols: List[str],
-            sample_weight_col: Optional[str]
+            sample_weight_col: Optional[str],
+            statement_params: Dict[str, str]
         ) -> float:
             import cloudpickle as cp
             import numpy as np
@@ -1171,14 +1194,14 @@ class AffinityPropagation(BaseTransformer):
             api_calls=[Session.call],
             custom_tags=dict([("autogen", True)]),
         )
-        score = session.call(
-            score_sproc_name,
+        score = score_wrapper_sproc(
+            session,
             query,
             stage_score_file_name,
             identifier.get_unescaped_names(self.input_cols),
             identifier.get_unescaped_names(self.label_cols),
             identifier.get_unescaped_names(self.sample_weight_col),
-            statement_params=statement_params,
+            statement_params,
         )
 
         cleanup_temp_files([local_score_file_name])
@@ -1196,18 +1219,20 @@ class AffinityPropagation(BaseTransformer):
             if self._sklearn_object._estimator_type == 'classifier':
                 outputs = _infer_signature(dataset[self.label_cols], "output")  # label columns is the desired type for output
                 outputs = _rename_features(outputs, self.output_cols)  # rename the output columns
-                self._model_signature_dict["predict"] = ModelSignature(inputs, outputs)
+                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, outputs)
-
+                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}_"
                 output_column_names = self._get_output_column_names(output_cols_prefix)
                 outputs = [FeatureSpec(dtype=DataType.DOUBLE, name=c) for c in output_column_names]
-                self._model_signature_dict[prob_func] = ModelSignature(inputs, outputs)
+                self._model_signature_dict[prob_func] = ModelSignature(inputs, 
+                                                                       ([] if self._drop_input_cols else inputs) + outputs)
 
     @property
     def model_signatures(self) -> Dict[str, ModelSignature]: