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

snowflake/ml/model/type_hints.py

@@ -4,17 +4,18 @@ from typing import TYPE_CHECKING, Sequence, TypedDict, TypeVar, Union
 import numpy.typing as npt
 from typing_extensions import NotRequired, TypeAlias
 
-from snowflake.ml.modeling.framework import base
-
 if TYPE_CHECKING:
     import numpy as np
     import pandas as pd
     import sklearn.base
     import sklearn.pipeline
+    import tensorflow
+    import torch
     import xgboost
 
     import snowflake.ml.model.custom_model
     import snowflake.snowpark
+    from snowflake.ml.modeling.framework import base  # noqa: F401
 
 
 _SupportedBuiltins = Union[int, float, bool, str, bytes, "_SupportedBuiltinsList"]
@@ -35,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"]
@@ -52,9 +54,11 @@ SupportedLocalModelType = Union[
     "sklearn.pipeline.Pipeline",
     "xgboost.XGBModel",
     "xgboost.Booster",
+    "torch.nn.Module",
+    "torch.jit.ScriptModule",  # type:ignore[name-defined]
 ]
 
-SupportedSnowMLModelType: TypeAlias = base.BaseEstimator
+SupportedSnowMLModelType: TypeAlias = "base.BaseEstimator"
 
 SupportedModelType = Union[
     SupportedLocalModelType,
@@ -71,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    |
 """
 
 
@@ -80,26 +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.
-
-    Internal-only options
-    _use_local_snowml: Use local SnowML when as the execution library of the deployment. If set to True, local SnowML
-        would be packed and uploaded to 1) session stage, if it is a temporary deployment, or 2) the provided stage path
-        if it is a permanent deployment. It should be set to True before SnowML available in Snowflake Anaconda Channel.
-        Default to False.
+    output_with_input_features: Whether or not preserve the input columns in the output when predicting.
+        Defaults to False.
     """
 
-    _use_local_snowml: NotRequired[bool]
-    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.
@@ -115,14 +118,16 @@ class WarehouseDeployOptions(DeployOptions):
 class ModelSaveOption(TypedDict):
     """Options for saving the model.
 
+    embed_local_ml_library: Embedding local SnowML into the code directory of the folder.
     allow_overwritten_stage_file: Flag to indicate when saving the model as a stage file, whether overwriting existed
         file is allowed. Default to False.
     """
 
+    embed_local_ml_library: NotRequired[bool]
     allow_overwritten_stage_file: NotRequired[bool]
 
 
-class CustomModelSaveOption(TypedDict):
+class CustomModelSaveOption(ModelSaveOption):
     ...
 
 
@@ -136,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)
@@ -682,26 +695,37 @@ class CalibratedClassifierCV(BaseTransformer):
         # input cols need to match unquoted / quoted
         input_cols = self.input_cols
         unquoted_input_cols = identifier.get_unescaped_names(self.input_cols)
+        quoted_input_cols = identifier.get_escaped_names(unquoted_input_cols)
 
         estimator = self._sklearn_object
 
-        input_df = dataset[input_cols] # Select input columns with quoted column names.
-        if hasattr(estimator, "feature_names_in_"):
-            missing_features = []
-            for i, f in enumerate(getattr(estimator, "feature_names_in_")):
-                if i >= len(input_cols) or (input_cols[i] != f and unquoted_input_cols[i] != f):
-                    missing_features.append(f)
-
-            if len(missing_features) > 0:
-                raise ValueError(
-                    "The feature names should match with those that were passed during fit.\n"
-                    f"Features seen during fit call but not present in the input: {missing_features}\n"
-                    f"Features in the input dataframe : {input_cols}\n"
-                )
-            input_df.columns = getattr(estimator, "feature_names_in_")
-        else:
-            # Just rename the column names to unquoted identifiers.
-            input_df.columns = unquoted_input_cols # Replace the quoted columns identifier with unquoted column ids.
+        features_required_by_estimator =  getattr(estimator, "feature_names_in_") if hasattr(estimator, "feature_names_in_") else unquoted_input_cols
+        missing_features = []
+        features_in_dataset = set(dataset.columns)
+        columns_to_select = []
+        for i, f in enumerate(features_required_by_estimator):
+            if (
+                    i >= len(input_cols)
+                or (input_cols[i] != f and unquoted_input_cols[i] != f and quoted_input_cols[i] != f)
+                or (input_cols[i] not in features_in_dataset and unquoted_input_cols[i] not in features_in_dataset 
+                    and quoted_input_cols[i] not in features_in_dataset)
+                ):
+                missing_features.append(f)
+            elif input_cols[i] in features_in_dataset:
+                columns_to_select.append(input_cols[i])
+            elif unquoted_input_cols[i] in features_in_dataset:
+                columns_to_select.append(unquoted_input_cols[i])
+            else:
+                columns_to_select.append(quoted_input_cols[i])
+
+        if len(missing_features) > 0:
+            raise ValueError(
+                "The feature names should match with those that were passed during fit.\n"
+                f"Features seen during fit call but not present in the input: {missing_features}\n"
+                f"Features in the input dataframe : {input_cols}\n"
+            )
+        input_df = dataset[columns_to_select]
+        input_df.columns = features_required_by_estimator
 
         transformed_numpy_array = getattr(estimator, inference_method)(
             input_df
@@ -782,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(
@@ -857,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):
@@ -1089,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(
@@ -1103,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,
@@ -1130,6 +1162,7 @@ class CalibratedClassifierCV(BaseTransformer):
             replace=True,
             session=session,
             statement_params=statement_params,
+            anonymous=True
         )
         def score_wrapper_sproc(
             session: Session,
@@ -1137,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
@@ -1187,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])
@@ -1212,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]: