snowflake-ml-python 1.6.0__py3-none-any.whl → 1.6.2__py3-none-any.whl

snowflake/ml/model/_packager/model_handlers/lightgbm.py

@@ -19,7 +19,11 @@ from typing_extensions import TypeGuard, Unpack
 from snowflake.ml._internal import type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
+from snowflake.ml.model._packager.model_handlers import (
+    _base,
+    _utils as handlers_utils,
+    model_objective_utils,
+)
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
@@ -41,47 +45,8 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model.pkl"
+    MODEL_BLOB_FILE_OR_DIR = "model.pkl"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
-    _BINARY_CLASSIFICATION_OBJECTIVES = ["binary"]
-    _MULTI_CLASSIFICATION_OBJECTIVES = ["multiclass", "multiclassova"]
-    _RANKING_OBJECTIVES = ["lambdarank", "rank_xendcg"]
-    _REGRESSION_OBJECTIVES = [
-        "regression",
-        "regression_l1",
-        "huber",
-        "fair",
-        "poisson",
-        "quantile",
-        "tweedie",
-        "mape",
-        "gamma",
-    ]
-
-    @classmethod
-    def get_model_objective(cls, model: Union["lightgbm.Booster", "lightgbm.LGBMModel"]) -> _base.ModelObjective:
-        import lightgbm
-
-        # does not account for cross-entropy and custom
-        if isinstance(model, lightgbm.LGBMClassifier):
-            num_classes = handlers_utils.get_num_classes_if_exists(model)
-            if num_classes == 2:
-                return _base.ModelObjective.BINARY_CLASSIFICATION
-            return _base.ModelObjective.MULTI_CLASSIFICATION
-        if isinstance(model, lightgbm.LGBMRanker):
-            return _base.ModelObjective.RANKING
-        if isinstance(model, lightgbm.LGBMRegressor):
-            return _base.ModelObjective.REGRESSION
-        model_objective = model.params["objective"]
-        if model_objective in cls._BINARY_CLASSIFICATION_OBJECTIVES:
-            return _base.ModelObjective.BINARY_CLASSIFICATION
-        if model_objective in cls._MULTI_CLASSIFICATION_OBJECTIVES:
-            return _base.ModelObjective.MULTI_CLASSIFICATION
-        if model_objective in cls._RANKING_OBJECTIVES:
-            return _base.ModelObjective.RANKING
-        if model_objective in cls._REGRESSION_OBJECTIVES:
-            return _base.ModelObjective.REGRESSION
-        return _base.ModelObjective.UNKNOWN
 
     @classmethod
     def can_handle(
@@ -116,6 +81,8 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.LGBMModelSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", True)
+
         import lightgbm
 
         assert isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel)
@@ -144,24 +111,25 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            if kwargs.get("enable_explainability", False):
-                output_type = model_signature.DataType.DOUBLE
-                if cls.get_model_objective(model) in [
-                    _base.ModelObjective.BINARY_CLASSIFICATION,
-                    _base.ModelObjective.MULTI_CLASSIFICATION,
-                ]:
-                    output_type = model_signature.DataType.STRING
+            model_objective_and_output = model_objective_utils.get_model_objective_and_output_type(model)
+            model_meta.model_objective = handlers_utils.validate_model_objective(
+                model_meta.model_objective, model_objective_and_output.objective
+            )
+            if enable_explainability:
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
                     explain_method="explain",
                     target_method="predict",
-                    output_return_type=output_type,
+                    output_return_type=model_objective_and_output.output_type,
                 )
+                model_meta.function_properties = {
+                    "explain": {model_meta_schema.FunctionProperties.PARTITIONED.value: False}
+                }
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
 
-        model_save_path = os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR)
+        model_save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
         with open(model_save_path, "wb") as f:
             cloudpickle.dump(model, f)
 
@@ -169,7 +137,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.LightGBMModelBlobOptions({"lightgbm_estimator_type": model.__class__.__name__}),
         )
         model_meta.models[name] = base_meta
@@ -182,11 +150,9 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
             ],
             check_local_version=True,
         )
-        if kwargs.get("enable_explainability", False):
-            model_meta.env.include_if_absent(
-                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
-                check_local_version=True,
-            )
+        if enable_explainability:
+            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
+            model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
 
         return None
 
@@ -226,6 +192,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         cls,
         raw_model: Union["lightgbm.Booster", "lightgbm.XGBModel"],
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.LGBMModelLoadOptions],
     ) -> custom_model.CustomModel:
         import lightgbm

snowflake/ml/model/_packager/model_handlers/llm.py

@@ -28,7 +28,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model"
+    MODEL_BLOB_FILE_OR_DIR = "model"
     LLM_META = "llm_meta"
     IS_AUTO_SIGNATURE = True
 
@@ -59,9 +59,12 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
         **kwargs: Unpack[model_types.LLMSaveOptions],
     ) -> None:
         assert not is_sub_model, "LLM can not be sub-model."
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for llm model.")
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
-        model_blob_dir_path = os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR)
+        model_blob_dir_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
 
         sig = model_signature.ModelSignature(
             inputs=[
@@ -86,7 +89,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.LLMModelBlobOptions(
                 {
                     "batch_size": model.max_batch_size,
@@ -143,6 +146,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
         cls,
         raw_model: llm.LLM,
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.LLMLoadOptions],
     ) -> custom_model.CustomModel:
         import gc
@@ -201,7 +205,9 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
                     "token": raw_model.token,
                 }
                 model_dir_path = raw_model.model_id_or_path
-                peft_config = peft.PeftConfig.from_pretrained(model_dir_path)  # type: ignore[attr-defined]
+                peft_config = peft.PeftConfig.from_pretrained(  # type: ignore[no-untyped-call, attr-defined]
+                    model_dir_path
+                )
                 base_model_path = peft_config.base_model_name_or_path
                 tokenizer = transformers.AutoTokenizer.from_pretrained(
                     base_model_path,
@@ -217,7 +223,7 @@ class LLMHandler(_base.BaseModelHandler[llm.LLM]):
                     model_dir_path,
                     device_map="auto",
                     torch_dtype="auto",
-                    **hub_kwargs,
+                    **hub_kwargs,  # type: ignore[arg-type]
                 )
                 hf_model.eval()
                 hf_model = hf_model.merge_and_unload()

snowflake/ml/model/_packager/model_handlers/mlflow.py

@@ -63,7 +63,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model"
+    MODEL_BLOB_FILE_OR_DIR = "model"
     _DEFAULT_TARGET_METHOD = "predict"
     DEFAULT_TARGET_METHODS = [_DEFAULT_TARGET_METHOD]
     IS_AUTO_SIGNATURE = True
@@ -97,6 +97,10 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.MLFlowSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for MLFlow model.")
+
         import mlflow
 
         assert isinstance(model, mlflow.pyfunc.PyFuncModel)
@@ -142,13 +146,13 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
             except (mlflow.MlflowException, OSError):
                 raise ValueError("Cannot load MLFlow model artifacts.")
 
-            file_utils.copy_file_or_tree(local_path, os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR))
+            file_utils.copy_file_or_tree(local_path, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
 
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.MLFlowModelBlobOptions({"artifact_path": model_info.artifact_path}),
         )
         model_meta.models[name] = base_meta
@@ -194,6 +198,7 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         cls,
         raw_model: "mlflow.pyfunc.PyFuncModel",
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.MLFlowLoadOptions],
     ) -> custom_model.CustomModel:
         from snowflake.ml.model import custom_model

snowflake/ml/model/_packager/model_handlers/model_objective_utils.py

@@ -0,0 +1,116 @@
+import json
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Union
+
+from snowflake.ml.model import model_signature, type_hints
+from snowflake.ml.model._packager.model_handlers import _utils as handlers_utils
+
+if TYPE_CHECKING:
+    import lightgbm
+    import xgboost
+
+
+@dataclass
+class ModelObjectiveAndOutputType:
+    objective: type_hints.ModelObjective
+    output_type: model_signature.DataType
+
+
+def get_model_objective_lightgbm(model: Union["lightgbm.Booster", "lightgbm.LGBMModel"]) -> type_hints.ModelObjective:
+
+    import lightgbm
+
+    _BINARY_CLASSIFICATION_OBJECTIVES = ["binary"]
+    _MULTI_CLASSIFICATION_OBJECTIVES = ["multiclass", "multiclassova"]
+    _RANKING_OBJECTIVES = ["lambdarank", "rank_xendcg"]
+    _REGRESSION_OBJECTIVES = [
+        "regression",
+        "regression_l1",
+        "huber",
+        "fair",
+        "poisson",
+        "quantile",
+        "tweedie",
+        "mape",
+        "gamma",
+    ]
+
+    # does not account for cross-entropy and custom
+    if isinstance(model, lightgbm.LGBMClassifier):
+        num_classes = handlers_utils.get_num_classes_if_exists(model)
+        if num_classes == 2:
+            return type_hints.ModelObjective.BINARY_CLASSIFICATION
+        return type_hints.ModelObjective.MULTI_CLASSIFICATION
+    if isinstance(model, lightgbm.LGBMRanker):
+        return type_hints.ModelObjective.RANKING
+    if isinstance(model, lightgbm.LGBMRegressor):
+        return type_hints.ModelObjective.REGRESSION
+    model_objective = model.params["objective"]
+    if model_objective in _BINARY_CLASSIFICATION_OBJECTIVES:
+        return type_hints.ModelObjective.BINARY_CLASSIFICATION
+    if model_objective in _MULTI_CLASSIFICATION_OBJECTIVES:
+        return type_hints.ModelObjective.MULTI_CLASSIFICATION
+    if model_objective in _RANKING_OBJECTIVES:
+        return type_hints.ModelObjective.RANKING
+    if model_objective in _REGRESSION_OBJECTIVES:
+        return type_hints.ModelObjective.REGRESSION
+    return type_hints.ModelObjective.UNKNOWN
+
+
+def get_model_objective_xgb(model: Union["xgboost.Booster", "xgboost.XGBModel"]) -> type_hints.ModelObjective:
+
+    import xgboost
+
+    _BINARY_CLASSIFICATION_OBJECTIVE_PREFIX = ["binary:"]
+    _MULTI_CLASSIFICATION_OBJECTIVE_PREFIX = ["multi:"]
+    _RANKING_OBJECTIVE_PREFIX = ["rank:"]
+    _REGRESSION_OBJECTIVE_PREFIX = ["reg:"]
+
+    model_objective = ""
+    if isinstance(model, xgboost.Booster):
+        model_params = json.loads(model.save_config())
+        model_objective = model_params.get("learner", {}).get("objective", "")
+    else:
+        if hasattr(model, "get_params"):
+            model_objective = model.get_params().get("objective", "")
+
+    if isinstance(model_objective, dict):
+        model_objective = model_objective.get("name", "")
+    for classification_objective in _BINARY_CLASSIFICATION_OBJECTIVE_PREFIX:
+        if classification_objective in model_objective:
+            return type_hints.ModelObjective.BINARY_CLASSIFICATION
+    for classification_objective in _MULTI_CLASSIFICATION_OBJECTIVE_PREFIX:
+        if classification_objective in model_objective:
+            return type_hints.ModelObjective.MULTI_CLASSIFICATION
+    for ranking_objective in _RANKING_OBJECTIVE_PREFIX:
+        if ranking_objective in model_objective:
+            return type_hints.ModelObjective.RANKING
+    for regression_objective in _REGRESSION_OBJECTIVE_PREFIX:
+        if regression_objective in model_objective:
+            return type_hints.ModelObjective.REGRESSION
+    return type_hints.ModelObjective.UNKNOWN
+
+
+def get_model_objective_and_output_type(model: Any) -> ModelObjectiveAndOutputType:
+    import xgboost
+
+    if isinstance(model, xgboost.Booster) or isinstance(model, xgboost.XGBModel):
+        model_objective = get_model_objective_xgb(model)
+        output_type = model_signature.DataType.DOUBLE
+        if model_objective == type_hints.ModelObjective.MULTI_CLASSIFICATION:
+            output_type = model_signature.DataType.STRING
+        return ModelObjectiveAndOutputType(objective=model_objective, output_type=output_type)
+
+    import lightgbm
+
+    if isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel):
+        model_objective = get_model_objective_lightgbm(model)
+        output_type = model_signature.DataType.DOUBLE
+        if model_objective in [
+            type_hints.ModelObjective.BINARY_CLASSIFICATION,
+            type_hints.ModelObjective.MULTI_CLASSIFICATION,
+        ]:
+            output_type = model_signature.DataType.STRING
+        return ModelObjectiveAndOutputType(objective=model_objective, output_type=output_type)
+
+    raise ValueError(f"Model type {type(model)} is not supported")

snowflake/ml/model/_packager/model_handlers/pytorch.py

@@ -37,7 +37,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model.pt"
+    MODEL_BLOB_FILE_OR_DIR = "model.pt"
     DEFAULT_TARGET_METHODS = ["forward"]
 
     @classmethod
@@ -73,6 +73,10 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.PyTorchSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for PyTorch model.")
+
         import torch
 
         assert isinstance(model, torch.nn.Module)
@@ -115,13 +119,13 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         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, cls.MODELE_BLOB_FILE_OR_DIR), "wb") as f:
+        with open(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
             torch.save(model, f, pickle_module=cloudpickle)
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -156,6 +160,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
         cls,
         raw_model: "torch.nn.Module",
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.PyTorchLoadOptions],
     ) -> custom_model.CustomModel:
         import torch

snowflake/ml/model/_packager/model_handlers/sentence_transformers.py

@@ -31,7 +31,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model"
+    MODEL_BLOB_FILE_OR_DIR = "model"
     DEFAULT_TARGET_METHODS = ["encode"]
 
     @classmethod
@@ -64,6 +64,10 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.SentenceTransformersSaveOptions],  # registry.log_model(options={...})
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for Sentence Transformer model.")
+
         # Validate target methods and signature (if possible)
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -101,14 +105,14 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         # save model
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
-        model.save(os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR))
+        model.save(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
 
         # save model metadata
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -154,6 +158,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         cls,
         raw_model: "sentence_transformers.SentenceTransformer",
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.SentenceTransformersLoadOptions],
     ) -> custom_model.CustomModel:
         import sentence_transformers

snowflake/ml/model/_packager/model_handlers/sklearn.py

@@ -6,6 +6,7 @@ import numpy as np
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
+import snowflake.snowpark.dataframe as sp_df
 from snowflake.ml._internal import type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
@@ -14,8 +15,13 @@ from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
+    model_meta_schema,
+)
+from snowflake.ml.model._signatures import (
+    numpy_handler,
+    snowpark_handler,
+    utils as model_signature_utils,
 )
-from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
     import sklearn.base
@@ -36,6 +42,27 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
 
     DEFAULT_TARGET_METHODS = ["predict", "transform", "predict_proba", "predict_log_proba", "decision_function"]
 
+    @classmethod
+    def get_model_objective(
+        cls, model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]
+    ) -> model_types.ModelObjective:
+        import sklearn.pipeline
+        from sklearn.base import is_classifier, is_regressor
+
+        if isinstance(model, sklearn.pipeline.Pipeline):
+            return model_types.ModelObjective.UNKNOWN
+        if is_regressor(model):
+            return model_types.ModelObjective.REGRESSION
+        if is_classifier(model):
+            classes_list = getattr(model, "classes_", [])
+            num_classes = getattr(model, "n_classes_", None) or len(classes_list)
+            if isinstance(num_classes, int):
+                if num_classes > 2:
+                    return model_types.ModelObjective.MULTI_CLASSIFICATION
+                return model_types.ModelObjective.BINARY_CLASSIFICATION
+            return model_types.ModelObjective.UNKNOWN
+        return model_types.ModelObjective.UNKNOWN
+
     @classmethod
     def can_handle(
         cls,
@@ -68,6 +95,18 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
 
         return cast(Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"], model)
 
+    @staticmethod
+    def get_explainability_supported_background(
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
+    ) -> Optional[pd.DataFrame]:
+        if isinstance(sample_input_data, pd.DataFrame) or isinstance(sample_input_data, sp_df.DataFrame):
+            return (
+                sample_input_data
+                if isinstance(sample_input_data, pd.DataFrame)
+                else snowpark_handler.SnowparkDataFrameHandler.convert_to_df(sample_input_data)
+            )
+        return None
+
     @classmethod
     def save_model(
         cls,
@@ -79,11 +118,31 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.SKLModelSaveOptions],
     ) -> None:
+        # setting None by default to distinguish if users did not set it
+        enable_explainability = kwargs.get("enable_explainability", None)
+
         import sklearn.base
         import sklearn.pipeline
 
         assert isinstance(model, sklearn.base.BaseEstimator) or isinstance(model, sklearn.pipeline.Pipeline)
 
+        background_data = cls.get_explainability_supported_background(sample_input_data)
+
+        # if users did not ask then we enable if we have background data
+        if enable_explainability is None and background_data is not None:
+            enable_explainability = True
+        if enable_explainability:
+            # if users set it explicitly but no background data then error out
+            if background_data is None:
+                raise ValueError(
+                    "Sample input data is required to enable explainability. Currently we only support this for "
+                    + "`pandas.DataFrame` and `snowflake.snowpark.dataframe.DataFrame`."
+                )
+            data_blob_path = os.path.join(model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR)
+            os.makedirs(data_blob_path, exist_ok=True)
+            with open(os.path.join(data_blob_path, name + cls.BG_DATA_FILE_SUFFIX), "wb") as f:
+                background_data.to_parquet(f)
+
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
                 model=model,
@@ -110,19 +169,38 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 get_prediction_fn=get_prediction,
             )
 
+            model_objective = cls.get_model_objective(model)
+            model_meta.model_objective = model_objective
+
+            if enable_explainability:
+                output_type = model_signature.DataType.DOUBLE
+
+                if model_objective == model_types.ModelObjective.MULTI_CLASSIFICATION:
+                    output_type = model_signature.DataType.STRING
+                model_meta = handlers_utils.add_explain_method_signature(
+                    model_meta=model_meta,
+                    explain_method="explain",
+                    target_method="predict",
+                    output_return_type=output_type,
+                )
+
         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, cls.MODELE_BLOB_FILE_OR_DIR), "wb") as f:
+        with open(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
             cloudpickle.dump(model, f)
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
+        if enable_explainability:
+            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
+            model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
+
         model_meta.env.include_if_absent(
             [model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn")], check_local_version=True
         )
@@ -153,6 +231,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         cls,
         raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.SKLModelLoadOptions],
     ) -> custom_model.CustomModel:
         from snowflake.ml.model import custom_model
@@ -165,6 +244,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
                 signature: model_signature.ModelSignature,
                 target_method: str,
+                background_data: Optional[pd.DataFrame],
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
@@ -179,11 +259,26 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
 
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
+                @custom_model.inference_api
+                def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    import shap
+
+                    # TODO: if not resolved by explainer, we need to pass the callable function
+                    try:
+                        explainer = shap.Explainer(raw_model, background_data)
+                        df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
+                    except TypeError as e:
+                        raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                if target_method == "explain":
+                    return explain_fn
+
                 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)
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name, background_data)
 
             _SKLModel = type(
                 "_SKLModel",