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

snowflake/ml/model/_handlers/pytorch.py

@@ -0,0 +1,186 @@
+import os
+import sys
+from typing import TYPE_CHECKING, Callable, Optional, Type, cast
+
+import cloudpickle
+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 _PyTorchHandler(_base._ModelHandler["torch.nn.Module"]):
+    """Handler for PyTorch based model.
+
+    Currently torch.nn.Module based classes are supported.
+    """
+
+    handler_type = "pytorch"
+    MODEL_BLOB_FILE = "model.pt"
+    DEFAULT_TARGET_METHODS = ["forward"]
+
+    @staticmethod
+    def can_handle(
+        model: model_types.SupportedModelType,
+    ) -> TypeGuard["torch.nn.Module"]:
+        return type_utils.LazyType("torch.nn.Module").isinstance(model) and not type_utils.LazyType(
+            "torch.jit.ScriptModule"
+        ).isinstance(model)
+
+    @staticmethod
+    def cast_model(
+        model: model_types.SupportedModelType,
+    ) -> "torch.nn.Module":
+        import torch
+
+        assert isinstance(model, torch.nn.Module)
+
+        return cast(torch.nn.Module, model)
+
+    @staticmethod
+    def _save_model(
+        name: str,
+        model: "torch.nn.Module",
+        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.PyTorchSaveOptions],
+    ) -> None:
+        import torch
+
+        assert isinstance(model, torch.nn.Module)
+
+        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=_PyTorchHandler.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,
+            )
+
+        # Torch.save using pickle will not pickle the model definition if defined in the top level of a module.
+        # Make sure that the module where the model is defined get pickled by value as well.
+        cloudpickle.register_pickle_by_value(sys.modules[model.__module__])
+        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, _PyTorchHandler.MODEL_BLOB_FILE), "wb") as f:
+            torch.save(model, f, pickle_module=cloudpickle)
+        base_meta = model_meta_api._ModelBlobMetadata(
+            name=name, model_type=_PyTorchHandler.handler_type, path=_PyTorchHandler.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.nn.Module":
+        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.load(f)
+        assert isinstance(m, torch.nn.Module)
+        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.
+        """
+        import torch
+
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: "torch.nn.Module",
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: "torch.nn.Module",
+                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.")
+
+                    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)
+
+            _PyTorchModel = type(
+                "_PyTorchModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _PyTorchModel
+
+        raw_model = _PyTorchHandler._load_model(name, model_meta, model_blobs_dir_path)
+        _PyTorchModel = _create_custom_model(raw_model, model_meta)
+        pytorch_model = _PyTorchModel(custom_model.ModelContext())
+
+        return pytorch_model

snowflake/ml/model/_handlers/sklearn.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Optional, Sequence, Type, Union, cast
+from typing import TYPE_CHECKING, Callable, Optional, Type, Union, cast
 
 import cloudpickle
 import numpy as np
@@ -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
@@ -80,6 +81,9 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
             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)
@@ -101,6 +105,7 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
             name=name, model_type=_SKLModelHandler.handler_type, path=_SKLModelHandler.MODEL_BLOB_FILE
         )
         model_meta.models[name] = base_meta
+        model_meta._include_if_absent([model_meta_api.Dependency(conda_name="scikit-learn", pip_name="scikit-learn")])
 
     @staticmethod
     def _load_model(
@@ -146,7 +151,7 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
-                output_col_names: Sequence[str],
+                signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
@@ -155,17 +160,18 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
 
                     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)
 
             _SKLModel = type(
                 "_SKLModel",

snowflake/ml/model/_handlers/snowmlmodel.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Optional, Sequence, Type, cast
+from typing import TYPE_CHECKING, Callable, Optional, Type, cast
 
 import cloudpickle
 import numpy as np
@@ -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
@@ -81,6 +82,9 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
                 def get_prediction(
                     target_method_name: str, sample_input: model_types.SupportedLocalDataType
                 ) -> model_types.SupportedLocalDataType:
+                    if not isinstance(sample_input, (pd.DataFrame,)):
+                        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)
@@ -106,7 +110,7 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
         model_dependencies = model._get_dependencies()
         for dep in model_dependencies:
             pkg_name = dep.split("==")[0]
-            _include_if_absent_pkgs.append((pkg_name, pkg_name))
+            _include_if_absent_pkgs.append(model_meta_api.Dependency(conda_name=pkg_name, pip_name=pkg_name))
         model_meta._include_if_absent(_include_if_absent_pkgs)
 
     @staticmethod
@@ -150,7 +154,7 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "BaseEstimator",
-                output_col_names: Sequence[str],
+                signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
@@ -159,17 +163,18 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
 
                     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)
 
             _SnowMLModel = type(
                 "_SnowMLModel",

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([("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,10 +1,11 @@
 import os
+import posixpath
 import tempfile
 import warnings
 from types import ModuleType
-from typing import Dict, List, Literal, Optional, Tuple, Union, overload
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, overload
 
-from snowflake.ml._internal import file_utils
+from snowflake.ml._internal import file_utils, type_utils
 from snowflake.ml.model import (
     _env,
     _model_handler,
@@ -13,9 +14,11 @@ from snowflake.ml.model import (
     model_signature,
     type_hints as model_types,
 )
-from snowflake.ml.modeling.framework import base
 from snowflake.snowpark import FileOperation, Session
 
+if TYPE_CHECKING:
+    from snowflake.ml.modeling.framework import base
+
 MODEL_BLOBS_DIR = "models"
 
 
@@ -23,7 +26,7 @@ MODEL_BLOBS_DIR = "models"
 def save_model(
     *,
     name: str,
-    model: base.BaseEstimator,
+    model: "base.BaseEstimator",
     model_dir_path: str,
     metadata: Optional[Dict[str, str]] = None,
     conda_dependencies: Optional[List[str]] = None,
@@ -135,7 +138,7 @@ def save_model(
 def save_model(
     *,
     name: str,
-    model: base.BaseEstimator,
+    model: "base.BaseEstimator",
     session: Session,
     model_stage_file_path: str,
     metadata: Optional[Dict[str, str]] = None,
@@ -322,9 +325,11 @@ def save_model(
             + f"{'None' if model_stage_file_path is None else 'specified'} at the same time."
         )
 
-    if ((signatures is None) and (sample_input is None) and not isinstance(model, base.BaseEstimator)) or (
-        (signatures is not None) and (sample_input is not None)
-    ):
+    if (
+        (signatures is None)
+        and (sample_input is None)
+        and not type_utils.LazyType("snowflake.ml.modeling.framework.base.BaseEstimator").isinstance(model)
+    ) or ((signatures is not None) and (sample_input is not None)):
         raise ValueError(
             "Signatures and sample_input both cannot be "
             + f"{'None for local model' if signatures is None else 'specified'} at the same time."
@@ -360,8 +365,8 @@ def save_model(
         )
 
     assert session and model_stage_file_path
-    if os.path.splitext(model_stage_file_path)[1] != ".zip":
-        raise ValueError("Provided model path in the stage {model_stage_file_path} must be a path to a zip file.")
+    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:
         meta = _save(
@@ -397,15 +402,15 @@ def _save(
     name: str,
     model: model_types.SupportedModelType,
     local_dir_path: str,
-    signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-    sample_input: Optional[model_types.SupportedDataType] = None,
-    metadata: Optional[Dict[str, str]] = None,
-    conda_dependencies: Optional[List[str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    python_version: Optional[str] = None,
-    ext_modules: Optional[List[ModuleType]] = None,
-    code_paths: Optional[List[str]] = None,
-    options: Optional[model_types.ModelSaveOption] = None,
+    signatures: Optional[Dict[str, model_signature.ModelSignature]],
+    sample_input: Optional[model_types.SupportedDataType],
+    metadata: Optional[Dict[str, str]],
+    conda_dependencies: Optional[List[str]],
+    pip_requirements: Optional[List[str]],
+    python_version: Optional[str],
+    ext_modules: Optional[List[ModuleType]],
+    code_paths: Optional[List[str]],
+    options: model_types.ModelSaveOption,
 ) -> _model_meta.ModelMetadata:
     local_dir_path = os.path.normpath(local_dir_path)
 
@@ -423,6 +428,7 @@ def _save(
         conda_dependencies=conda_dependencies,
         pip_requirements=pip_requirements,
         python_version=python_version,
+        **options,
     ) as meta:
         model_blobs_path = os.path.join(local_dir_path, MODEL_BLOBS_DIR)
         os.makedirs(model_blobs_path, exist_ok=True)
@@ -538,8 +544,8 @@ 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":
-        raise ValueError("Provided model path in the stage {model_stage_file_path} must be a path to a zip file.")
+    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)
     zf = fo.get_stream(model_stage_file_path)