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

snowflake/ml/model/_handlers/torchscript.py

@@ -0,0 +1,180 @@
+import os
+from typing import TYPE_CHECKING, Callable, Optional, Type, cast
+
+import pandas as pd
+from typing_extensions import TypeGuard, Unpack
+
+from snowflake.ml._internal import type_utils
+from snowflake.ml.model import (
+    _model_meta as model_meta_api,
+    custom_model,
+    model_signature,
+    type_hints as model_types,
+)
+from snowflake.ml.model._handlers import _base
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _TorchScriptHandler(_base._ModelHandler["torch.jit.ScriptModule"]):  # type:ignore[name-defined]
+    """Handler for PyTorch JIT based model.
+
+    Currently torch.jit.ScriptModule based classes are supported.
+    """
+
+    handler_type = "torchscript"
+    MODEL_BLOB_FILE = "model.pt"
+    DEFAULT_TARGET_METHODS = ["forward"]
+
+    @staticmethod
+    def can_handle(
+        model: model_types.SupportedModelType,
+    ) -> TypeGuard["torch.jit.ScriptModule"]:  # type:ignore[name-defined]
+        return type_utils.LazyType("torch.jit.ScriptModule").isinstance(model)
+
+    @staticmethod
+    def cast_model(
+        model: model_types.SupportedModelType,
+    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+        import torch
+
+        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+
+        return cast(torch.jit.ScriptModule, model)  # type:ignore[name-defined]
+
+    @staticmethod
+    def _save_model(
+        name: str,
+        model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        sample_input: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
+        **kwargs: Unpack[model_types.TorchScriptSaveOptions],
+    ) -> None:
+        import torch
+
+        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+
+        if not is_sub_model:
+            target_methods = model_meta_api._get_target_methods(
+                model=model,
+                target_methods=kwargs.pop("target_methods", None),
+                default_target_methods=_TorchScriptHandler.DEFAULT_TARGET_METHODS,
+            )
+
+            def get_prediction(
+                target_method_name: str, sample_input: "model_types.SupportedLocalDataType"
+            ) -> model_types.SupportedLocalDataType:
+                if not model_signature._SeqOfPyTorchTensorHandler.can_handle(sample_input):
+                    sample_input = model_signature._SeqOfPyTorchTensorHandler.convert_from_df(
+                        model_signature._convert_local_data_to_df(sample_input)
+                    )
+
+                model.eval()
+                target_method = getattr(model, target_method_name, None)
+                assert callable(target_method)
+                with torch.no_grad():
+                    predictions_df = target_method(sample_input)
+                return predictions_df
+
+            model_meta = model_meta_api._validate_signature(
+                model=model,
+                model_meta=model_meta,
+                target_methods=target_methods,
+                sample_input=sample_input,
+                get_prediction_fn=get_prediction,
+            )
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        os.makedirs(model_blob_path, exist_ok=True)
+        with open(os.path.join(model_blob_path, _TorchScriptHandler.MODEL_BLOB_FILE), "wb") as f:
+            torch.jit.save(model, f)  # type:ignore[attr-defined]
+        base_meta = model_meta_api._ModelBlobMetadata(
+            name=name, model_type=_TorchScriptHandler.handler_type, path=_TorchScriptHandler.MODEL_BLOB_FILE
+        )
+        model_meta.models[name] = base_meta
+        model_meta._include_if_absent([model_meta_api.Dependency(conda_name="pytorch", pip_name="torch")])
+
+    @staticmethod
+    def _load_model(
+        name: str, model_meta: model_meta_api.ModelMetadata, model_blobs_dir_path: str
+    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+        import torch
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        if not hasattr(model_meta, "models"):
+            raise ValueError("Ill model metadata found.")
+        model_blobs_metadata = model_meta.models
+        if name not in model_blobs_metadata:
+            raise ValueError(f"Blob of model {name} does not exist.")
+        model_blob_metadata = model_blobs_metadata[name]
+        model_blob_filename = model_blob_metadata.path
+        with open(os.path.join(model_blob_path, model_blob_filename), "rb") as f:
+            m = torch.jit.load(f)  # type:ignore[attr-defined]
+        assert isinstance(m, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        return m
+
+    @staticmethod
+    def _load_as_custom_model(
+        name: str, model_meta: model_meta_api.ModelMetadata, model_blobs_dir_path: str
+    ) -> custom_model.CustomModel:
+        """Create a custom model class wrap for unified interface when being deployed. The predict method will be
+        re-targeted based on target_method metadata.
+
+        Args:
+            name: Name of the model.
+            model_meta: The model metadata.
+            model_blobs_dir_path: Directory path to the whole model.
+
+        Returns:
+            The model object as a custom model.
+        """
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+                signature: model_signature.ModelSignature,
+                target_method: str,
+            ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
+                @custom_model.inference_api
+                def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    if X.isnull().any(axis=None):
+                        raise ValueError("Tensor cannot handle null values.")
+
+                    import torch
+
+                    raw_model.eval()
+
+                    t = model_signature._SeqOfPyTorchTensorHandler.convert_from_df(X, signature.inputs)
+
+                    with torch.no_grad():
+                        res = getattr(raw_model, target_method)(t)
+                    return model_signature._rename_pandas_df(
+                        data=model_signature._SeqOfPyTorchTensorHandler.convert_to_df(res), features=signature.outputs
+                    )
+
+                return fn
+
+            type_method_dict = {}
+            for target_method_name, sig in model_meta.signatures.items():
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
+
+            _TorchScriptModel = type(
+                "_TorchScriptModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _TorchScriptModel
+
+        raw_model = _TorchScriptHandler._load_model(name, model_meta, model_blobs_dir_path)
+        _TorchScriptModel = _create_custom_model(raw_model, model_meta)
+        torchscript_model = _TorchScriptModel(custom_model.ModelContext())
+
+        return torchscript_model

snowflake/ml/model/_handlers/xgboost.py

@@ -1,6 +1,6 @@
 # mypy: disable-error-code="import"
 import os
-from typing import TYPE_CHECKING, Callable, Optional, Sequence, Type, Union
+from typing import TYPE_CHECKING, Callable, Optional, Type, Union
 
 import numpy as np
 import pandas as pd
@@ -10,6 +10,7 @@ from snowflake.ml._internal import type_utils
 from snowflake.ml.model import (
     _model_meta as model_meta_api,
     custom_model,
+    model_signature,
     type_hints as model_types,
 )
 from snowflake.ml.model._handlers import _base
@@ -72,6 +73,9 @@ class _XGBModelHandler(_base._ModelHandler[Union["xgboost.Booster", "xgboost.XGB
             def get_prediction(
                 target_method_name: str, sample_input: model_types.SupportedLocalDataType
             ) -> model_types.SupportedLocalDataType:
+                if not isinstance(sample_input, (pd.DataFrame, np.ndarray)):
+                    sample_input = model_signature._convert_local_data_to_df(sample_input)
+
                 target_method = getattr(model, target_method_name, None)
                 assert callable(target_method)
                 predictions_df = target_method(sample_input)
@@ -95,7 +99,12 @@ class _XGBModelHandler(_base._ModelHandler[Union["xgboost.Booster", "xgboost.XGB
             options={"xgb_estimator_type": model.__class__.__name__},
         )
         model_meta.models[name] = base_meta
-        model_meta._include_if_absent([("scikit-learn", "scikit-learn"), ("xgboost", "xgboost")])
+        model_meta._include_if_absent(
+            [
+                model_meta_api.Dependency(conda_name="scikit-learn", pip_name="scikit-learn"),
+                model_meta_api.Dependency(conda_name="xgboost", pip_name="xgboost"),
+            ]
+        )
 
     @staticmethod
     def _load_model(
@@ -143,7 +152,7 @@ class _XGBModelHandler(_base._ModelHandler[Union["xgboost.Booster", "xgboost.XGB
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
-                output_col_names: Sequence[str],
+                signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
@@ -152,17 +161,18 @@ class _XGBModelHandler(_base._ModelHandler[Union["xgboost.Booster", "xgboost.XGB
 
                     if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
                         # In case of multi-output estimators, predict_proba(), decision_function(), etc., functions
-                        # return a list of ndarrays. We need to concatenate them.
-                        res = np.concatenate(res, axis=1)
-                    return pd.DataFrame(res, columns=output_col_names)
+                        # return a list of ndarrays. We need to deal them seperately
+                        df = model_signature._SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature._rename_pandas_df(df, signature.outputs)
 
                 return fn
 
             type_method_dict = {}
             for target_method_name, sig in model_meta.signatures.items():
-                type_method_dict[target_method_name] = fn_factory(
-                    raw_model, [spec.name for spec in sig.outputs], target_method_name
-                )
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 
             _XGBModel = type(
                 "_XGBModel",

snowflake/ml/model/_model.py

@@ -1,4 +1,5 @@
 import os
+import posixpath
 import tempfile
 import warnings
 from types import ModuleType
@@ -364,7 +365,7 @@ def save_model(
         )
 
     assert session and model_stage_file_path
-    if os.path.splitext(model_stage_file_path)[1] != ".zip":
+    if posixpath.splitext(model_stage_file_path)[1] != ".zip":
         raise ValueError(f"Provided model path in the stage {model_stage_file_path} must be a path to a zip file.")
 
     with tempfile.TemporaryDirectory() as temp_local_model_dir_path:
@@ -543,7 +544,7 @@ def load_model(
         return _load(local_dir_path=model_dir_path, meta_only=meta_only)
 
     assert session and model_stage_file_path
-    if os.path.splitext(model_stage_file_path)[1] != ".zip":
+    if posixpath.splitext(model_stage_file_path)[1] != ".zip":
         raise ValueError(f"Provided model path in the stage {model_stage_file_path} must be a path to a zip file.")
 
     fo = FileOperation(session=session)

snowflake/ml/model/_model_meta.py

@@ -3,10 +3,11 @@ import importlib
 import os
 import sys
 import warnings
+from collections import namedtuple
 from contextlib import contextmanager
 from datetime import datetime
 from types import ModuleType
-from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, cast
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, cast
 
 import cloudpickle
 import yaml
@@ -24,6 +25,8 @@ from snowflake.snowpark import DataFrame as SnowparkDataFrame
 MODEL_METADATA_VERSION = 1
 _BASIC_DEPENDENCIES = _core_requirements.REQUIREMENTS
 
+Dependency = namedtuple("Dependency", ["conda_name", "pip_name"])
+
 
 @dataclasses.dataclass
 class _ModelBlobMetadata:
@@ -214,9 +217,11 @@ class ModelMetadata:
             pip_requirements if pip_requirements else []
         )
         if "local_ml_library_version" in kwargs:
-            self._include_if_absent([(dep, dep) for dep in _BASIC_DEPENDENCIES])
+            self._include_if_absent([Dependency(conda_name=dep, pip_name=dep) for dep in _BASIC_DEPENDENCIES])
         else:
-            self._include_if_absent([(dep, dep) for dep in _BASIC_DEPENDENCIES + [env_utils._SNOWML_PKG_NAME]])
+            self._include_if_absent(
+                [Dependency(conda_name=dep, pip_name=dep) for dep in _BASIC_DEPENDENCIES + [env_utils._SNOWML_PKG_NAME]]
+            )
 
         self.__dict__.update(kwargs)
 
@@ -234,7 +239,7 @@ class ModelMetadata:
             for req in reqs
         )
 
-    def _include_if_absent(self, pkgs: List[Tuple[str, str]]) -> None:
+    def _include_if_absent(self, pkgs: List[Dependency]) -> None:
         conda_reqs_str, pip_reqs_str = tuple(zip(*pkgs))
         pip_reqs = env_utils.validate_pip_requirement_string_list(list(pip_reqs_str))
         conda_reqs = env_utils.validate_conda_dependency_string_list(list(conda_reqs_str))
@@ -327,7 +332,7 @@ class ModelMetadata:
             path: The path of the directory to write a yaml file in it.
         """
         model_yaml_path = os.path.join(path, ModelMetadata.MODEL_METADATA_FILE)
-        with open(model_yaml_path, "w") as out:
+        with open(model_yaml_path, "w", encoding="utf-8") as out:
             yaml.safe_dump({**self.to_dict(), "version": MODEL_METADATA_VERSION}, stream=out, default_flow_style=False)
 
         env_dir_path = os.path.join(path, ModelMetadata.ENV_DIR)
@@ -350,7 +355,7 @@ class ModelMetadata:
             Loaded model metadata object.
         """
         model_yaml_path = os.path.join(path, ModelMetadata.MODEL_METADATA_FILE)
-        with open(model_yaml_path) as f:
+        with open(model_yaml_path, encoding="utf-8") as f:
             loaded_mata = yaml.safe_load(f.read())
 
         loaded_mata_version = loaded_mata.pop("version", None)
@@ -392,7 +397,7 @@ def _validate_signature(
     if isinstance(sample_input, SnowparkDataFrame):
         # Added because of Any from missing stubs.
         trunc_sample_input = cast(SnowparkDataFrame, trunc_sample_input)
-        local_sample_input = trunc_sample_input.to_pandas()
+        local_sample_input = model_signature._SnowparkDataFrameHandler.convert_to_df(trunc_sample_input)
     else:
         local_sample_input = trunc_sample_input
     for target_method in target_methods: