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

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

@@ -2,23 +2,67 @@ import json
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Union
 
+from snowflake.ml._internal import type_utils
 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 catboost
     import lightgbm
+    import sklearn
+    import sklearn.pipeline
     import xgboost
 
 
 @dataclass
-class ModelObjectiveAndOutputType:
-    objective: type_hints.ModelObjective
+class ModelTaskAndOutputType:
+    task: type_hints.Task
     output_type: model_signature.DataType
 
 
-def get_model_objective_lightgbm(model: Union["lightgbm.Booster", "lightgbm.LGBMModel"]) -> type_hints.ModelObjective:
+def get_task_skl(model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]) -> type_hints.Task:
+    from sklearn.base import is_classifier, is_regressor
+
+    if type_utils.LazyType("sklearn.pipeline.Pipeline").isinstance(model):
+        return type_hints.Task.UNKNOWN
+    if is_regressor(model):
+        return type_hints.Task.TABULAR_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 type_hints.Task.TABULAR_MULTI_CLASSIFICATION
+            return type_hints.Task.TABULAR_BINARY_CLASSIFICATION
+        return type_hints.Task.UNKNOWN
+    return type_hints.Task.UNKNOWN
+
+
+def get_model_task_catboost(model: "catboost.CatBoost") -> type_hints.Task:
+    loss_function = None
+    if type_utils.LazyType("catboost.CatBoost").isinstance(model):
+        loss_function = model.get_all_params()["loss_function"]  # type: ignore[attr-defined]
+
+    if (type_utils.LazyType("catboost.CatBoostClassifier").isinstance(model)) or model._is_classification_objective(
+        loss_function
+    ):
+        num_classes = handlers_utils.get_num_classes_if_exists(model)
+        if num_classes == 0:
+            return type_hints.Task.UNKNOWN
+        if num_classes <= 2:
+            return type_hints.Task.TABULAR_BINARY_CLASSIFICATION
+        return type_hints.Task.TABULAR_MULTI_CLASSIFICATION
+    if (type_utils.LazyType("catboost.CatBoostRanker").isinstance(model)) or model._is_ranking_objective(loss_function):
+        return type_hints.Task.TABULAR_RANKING
+    if (type_utils.LazyType("catboost.CatBoostRegressor").isinstance(model)) or model._is_regression_objective(
+        loss_function
+    ):
+        return type_hints.Task.TABULAR_REGRESSION
 
-    import lightgbm
+    return type_hints.Task.UNKNOWN
+
+
+def get_model_task_lightgbm(model: Union["lightgbm.Booster", "lightgbm.LGBMModel"]) -> type_hints.Task:
 
     _BINARY_CLASSIFICATION_OBJECTIVES = ["binary"]
     _MULTI_CLASSIFICATION_OBJECTIVES = ["multiclass", "multiclassova"]
@@ -36,81 +80,90 @@ def get_model_objective_lightgbm(model: Union["lightgbm.Booster", "lightgbm.LGBM
     ]
 
     # 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
+    model_task = ""
+    if type_utils.LazyType("lightgbm.Booster").isinstance(model):
+        model_task = model.params["objective"]  # type: ignore[attr-defined]
+    elif hasattr(model, "objective_"):
+        model_task = model.objective_
+    if model_task in _BINARY_CLASSIFICATION_OBJECTIVES:
+        return type_hints.Task.TABULAR_BINARY_CLASSIFICATION
+    if model_task in _MULTI_CLASSIFICATION_OBJECTIVES:
+        return type_hints.Task.TABULAR_MULTI_CLASSIFICATION
+    if model_task in _RANKING_OBJECTIVES:
+        return type_hints.Task.TABULAR_RANKING
+    if model_task in _REGRESSION_OBJECTIVES:
+        return type_hints.Task.TABULAR_REGRESSION
+    return type_hints.Task.UNKNOWN
+
+
+def get_model_task_xgb(model: Union["xgboost.Booster", "xgboost.XGBModel"]) -> type_hints.Task:
 
     _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", "")
+    model_task = ""
+    if type_utils.LazyType("xgboost.Booster").isinstance(model):
+        model_params = json.loads(model.save_config())  # type: ignore[attr-defined]
+        model_task = model_params.get("learner", {}).get("objective", "")
     else:
         if hasattr(model, "get_params"):
-            model_objective = model.get_params().get("objective", "")
+            model_task = model.get_params().get("objective", "")
 
-    if isinstance(model_objective, dict):
-        model_objective = model_objective.get("name", "")
+    if isinstance(model_task, dict):
+        model_task = model_task.get("name", "")
     for classification_objective in _BINARY_CLASSIFICATION_OBJECTIVE_PREFIX:
-        if classification_objective in model_objective:
-            return type_hints.ModelObjective.BINARY_CLASSIFICATION
+        if classification_objective in model_task:
+            return type_hints.Task.TABULAR_BINARY_CLASSIFICATION
     for classification_objective in _MULTI_CLASSIFICATION_OBJECTIVE_PREFIX:
-        if classification_objective in model_objective:
-            return type_hints.ModelObjective.MULTI_CLASSIFICATION
+        if classification_objective in model_task:
+            return type_hints.Task.TABULAR_MULTI_CLASSIFICATION
     for ranking_objective in _RANKING_OBJECTIVE_PREFIX:
-        if ranking_objective in model_objective:
-            return type_hints.ModelObjective.RANKING
+        if ranking_objective in model_task:
+            return type_hints.Task.TABULAR_RANKING
     for regression_objective in _REGRESSION_OBJECTIVE_PREFIX:
-        if regression_objective in model_objective:
-            return type_hints.ModelObjective.REGRESSION
-    return type_hints.ModelObjective.UNKNOWN
+        if regression_objective in model_task:
+            return type_hints.Task.TABULAR_REGRESSION
+    return type_hints.Task.UNKNOWN
 
 
-def get_model_objective_and_output_type(model: Any) -> ModelObjectiveAndOutputType:
-    import xgboost
+def get_model_task_and_output_type(model: Any) -> ModelTaskAndOutputType:
+    if type_utils.LazyType("xgboost.Booster").isinstance(model) or type_utils.LazyType("xgboost.XGBModel").isinstance(
+        model
+    ):
+        task = get_model_task_xgb(model)
+        output_type = model_signature.DataType.DOUBLE
+        if task == type_hints.Task.TABULAR_MULTI_CLASSIFICATION:
+            output_type = model_signature.DataType.STRING
+        return ModelTaskAndOutputType(task=task, output_type=output_type)
 
-    if isinstance(model, xgboost.Booster) or isinstance(model, xgboost.XGBModel):
-        model_objective = get_model_objective_xgb(model)
+    if type_utils.LazyType("lightgbm.Booster").isinstance(model) or type_utils.LazyType(
+        "lightgbm.LGBMModel"
+    ).isinstance(model):
+        task = get_model_task_lightgbm(model)
         output_type = model_signature.DataType.DOUBLE
-        if model_objective == type_hints.ModelObjective.MULTI_CLASSIFICATION:
+        if task in [
+            type_hints.Task.TABULAR_BINARY_CLASSIFICATION,
+            type_hints.Task.TABULAR_MULTI_CLASSIFICATION,
+        ]:
             output_type = model_signature.DataType.STRING
-        return ModelObjectiveAndOutputType(objective=model_objective, output_type=output_type)
+        return ModelTaskAndOutputType(task=task, output_type=output_type)
 
-    import lightgbm
+    if type_utils.LazyType("catboost.CatBoost").isinstance(model):
+        task = get_model_task_catboost(model)
+        output_type = model_signature.DataType.DOUBLE
+        if task == type_hints.Task.TABULAR_MULTI_CLASSIFICATION:
+            output_type = model_signature.DataType.STRING
+        return ModelTaskAndOutputType(task=task, output_type=output_type)
 
-    if isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel):
-        model_objective = get_model_objective_lightgbm(model)
+    if type_utils.LazyType("sklearn.base.BaseEstimator").isinstance(model) or type_utils.LazyType(
+        "sklearn.pipeline.Pipeline"
+    ).isinstance(model):
+        task = get_task_skl(model)
         output_type = model_signature.DataType.DOUBLE
-        if model_objective in [
-            type_hints.ModelObjective.BINARY_CLASSIFICATION,
-            type_hints.ModelObjective.MULTI_CLASSIFICATION,
-        ]:
+        if task == type_hints.Task.TABULAR_MULTI_CLASSIFICATION:
             output_type = model_signature.DataType.STRING
-        return ModelObjectiveAndOutputType(objective=model_objective, output_type=output_type)
+        return ModelTaskAndOutputType(task=task, output_type=output_type)
 
     raise ValueError(f"Model type {type(model)} is not supported")

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

@@ -2,7 +2,6 @@ import logging
 import os
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
 
-import cloudpickle
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
@@ -120,9 +119,21 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         model_meta.env.include_if_absent(
             [
                 model_env.ModelDependency(requirement="sentence-transformers", pip_name="sentence-transformers"),
+                model_env.ModelDependency(requirement="transformers", pip_name="transformers"),
+                model_env.ModelDependency(requirement="pytorch", pip_name="torch"),
             ],
             check_local_version=True,
         )
+        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+
+    @staticmethod
+    def _get_device_config(**kwargs: Unpack[model_types.SentenceTransformersLoadOptions]) -> Optional[str]:
+        if kwargs.get("device", None) is not None:
+            return kwargs["device"]
+        elif kwargs.get("use_gpu", False):
+            return "cuda"
+
+        return None
 
     @classmethod
     def load_model(
@@ -144,13 +155,9 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         model_blob_filename = model_blob_metadata.path
         model_blob_file_or_dir_path = os.path.join(model_blob_path, model_blob_filename)
 
-        if os.path.isdir(model_blob_file_or_dir_path):  # if the saved model is a directory
-            model = sentence_transformers.SentenceTransformer(model_blob_file_or_dir_path)
-        else:
-            assert os.path.isfile(model_blob_file_or_dir_path)  # if the saved model is a file
-            with open(model_blob_file_or_dir_path, "rb") as f:
-                model = cloudpickle.load(f)
-            assert isinstance(model, sentence_transformers.SentenceTransformer)
+        model = sentence_transformers.SentenceTransformer(
+            model_blob_file_or_dir_path, device=cls._get_device_config(**kwargs)
+        )
         return model
 
     @classmethod

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

@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, Union, cast, final
 
 import cloudpickle
@@ -6,22 +7,21 @@ 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
-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,
     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
@@ -40,28 +40,14 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    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
+    DEFAULT_TARGET_METHODS = [
+        "predict",
+        "transform",
+        "predict_proba",
+        "predict_log_proba",
+        "decision_function",
+    ]
+    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
 
     @classmethod
     def can_handle(
@@ -95,18 +81,6 @@ 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,
@@ -125,23 +99,10 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         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 users set it explicitly but no sample_input_data then error out
+            if sample_input_data is None:
+                raise ValueError("Sample input data is required to enable explainability.")
 
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -151,7 +112,8 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             )
 
             def get_prediction(
-                target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
+                target_method_name: str,
+                sample_input_data: model_types.SupportedLocalDataType,
             ) -> model_types.SupportedLocalDataType:
                 if not isinstance(sample_input_data, (pd.DataFrame, np.ndarray)):
                     sample_input_data = model_signature._convert_local_data_to_df(sample_input_data)
@@ -169,19 +131,40 @@ 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
+            explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
 
+            background_data = handlers_utils.get_explainability_supported_background(
+                sample_input_data, model_meta, explain_target_method
+            )
+
+            model_task_and_output_type = model_objective_utils.get_model_task_and_output_type(model)
+            model_meta.task = model_task_and_output_type.task
+
+            # if users did not ask then we enable if we have background data
+            if enable_explainability is None:
+                if background_data is None:
+                    warnings.warn(
+                        "sample_input_data should be provided to enable explainability by default",
+                        category=UserWarning,
+                        stacklevel=1,
+                    )
+                    enable_explainability = False
+                else:
+                    enable_explainability = True
             if enable_explainability:
-                output_type = model_signature.DataType.DOUBLE
+                handlers_utils.save_background_data(
+                    model_blobs_dir_path,
+                    cls.EXPLAIN_ARTIFACTS_DIR,
+                    cls.BG_DATA_FILE_SUFFIX,
+                    name,
+                    background_data,
+                )
 
-                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,
+                    target_method=explain_target_method,
+                    output_return_type=model_task_and_output_type.output_type,
                 )
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
@@ -202,7 +185,8 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             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
+            [model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn")],
+            check_local_version=True,
         )
 
     @classmethod

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

@@ -43,6 +43,8 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     DEFAULT_TARGET_METHODS = ["predict", "transform", "predict_proba", "predict_log_proba", "decision_function"]
+    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
+
     IS_AUTO_SIGNATURE = True
 
     @classmethod
@@ -71,13 +73,14 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
 
     @classmethod
     def _get_local_version_package(cls, pkg_name: str) -> Optional[version.Version]:
-        import importlib_metadata
+        from importlib import metadata as importlib_metadata
+
         from packaging import version
 
         local_version = None
 
         try:
-            local_dist = importlib_metadata.distribution(pkg_name)  # type: ignore[no-untyped-call]
+            local_dist = importlib_metadata.distribution(pkg_name)
             local_version = version.parse(local_dist.version)
         except importlib_metadata.PackageNotFoundError:
             pass
@@ -104,7 +107,13 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     def _get_supported_object_for_explainability(
         cls, estimator: "BaseEstimator", enable_explainability: Optional[bool]
     ) -> Any:
-        methods = ["to_xgboost", "to_lightgbm"]
+        from snowflake.ml.modeling import pipeline as snowml_pipeline
+
+        # handle pipeline objects separately
+        if isinstance(estimator, snowml_pipeline.Pipeline):  # type: ignore[attr-defined]
+            return None
+
+        methods = ["to_xgboost", "to_lightgbm", "to_sklearn"]
         for method_name in methods:
             if hasattr(estimator, method_name):
                 try:
@@ -136,9 +145,9 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         # Pipeline is inherited from BaseEstimator, so no need to add one more check
 
         if not is_sub_model:
-            if sample_input_data is not None or model_meta.signatures:
+            if model_meta.signatures:
                 warnings.warn(
-                    "Inferring model signature from sample input or providing model signature for Snowpark ML "
+                    "Providing model signature for Snowpark ML "
                     + "Modeling model is not required. Model signature will automatically be inferred during fitting. ",
                     UserWarning,
                     stacklevel=2,
@@ -162,22 +171,31 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
             python_base_obj = cls._get_supported_object_for_explainability(model, enable_explainability)
             if python_base_obj is None:
                 if enable_explainability:  # if user set enable_explainability to True, throw error else silently skip
-                    raise ValueError("Explain only support for xgboost or lightgbm Snowpark ML models.")
+                    raise ValueError(
+                        "Explain only supported for xgboost, lightgbm and sklearn (not pipeline) Snowpark ML models."
+                    )
                 # set None to False so we don't include shap in the environment
                 enable_explainability = False
             else:
-                model_objective_and_output_type = model_objective_utils.get_model_objective_and_output_type(
-                    python_base_obj
-                )
-                model_meta.model_objective = model_objective_and_output_type.objective
+                model_task_and_output_type = model_objective_utils.get_model_task_and_output_type(python_base_obj)
+                model_meta.task = model_task_and_output_type.task
+                explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
                     explain_method="explain",
-                    target_method="predict",
-                    output_return_type=model_objective_and_output_type.output_type,
+                    target_method=explain_target_method,
+                    output_return_type=model_task_and_output_type.output_type,
                 )
                 enable_explainability = True
 
+                background_data = handlers_utils.get_explainability_supported_background(
+                    sample_input_data, model_meta, explain_target_method
+                )
+                if background_data is not None:
+                    handlers_utils.save_background_data(
+                        model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, cls.BG_DATA_FILE_SUFFIX, name, background_data
+                    )
+
         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.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
@@ -258,6 +276,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 raw_model: "BaseEstimator",
                 signature: model_signature.ModelSignature,
                 target_method: str,
+                background_data: Optional[pd.DataFrame] = None,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
@@ -276,16 +295,16 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
                     import shap
 
-                    methods = ["to_xgboost", "to_lightgbm"]
+                    methods = ["to_xgboost", "to_lightgbm", "to_sklearn"]
                     for method_name in methods:
                         try:
                             base_model = getattr(raw_model, method_name)()
-                            explainer = shap.TreeExplainer(base_model)
-                            df = pd.DataFrame(explainer(X).values)
+                            explainer = shap.Explainer(base_model, masker=background_data)
+                            df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
                             return model_signature_utils.rename_pandas_df(df, signature.outputs)
                         except exceptions.SnowflakeMLException:
                             pass  # Do nothing and continue to the next method
-                    raise ValueError("The model must be an xgboost or lightgbm estimator.")
+                    raise ValueError("The model must be an xgboost, lightgbm or sklearn (not pipeline) estimator.")
 
                 if target_method == "explain":
                     return explain_fn
@@ -294,7 +313,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
 
             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)
 
             _SnowMLModel = type(
                 "_SnowMLModel",

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

@@ -1,6 +1,7 @@
 # mypy: disable-error-code="import"
 import os
 import warnings
+from importlib import metadata as importlib_metadata
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -13,7 +14,6 @@ from typing import (
     final,
 )
 
-import importlib_metadata
 import numpy as np
 import pandas as pd
 from packaging import version
@@ -53,6 +53,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
 
     MODEL_BLOB_FILE_OR_DIR = "model.ubj"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
+    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba"]
 
     @classmethod
     def can_handle(
@@ -96,7 +97,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         local_xgb_version = None
 
         try:
-            local_dist = importlib_metadata.distribution("xgboost")  # type: ignore[no-untyped-call]
+            local_dist = importlib_metadata.distribution("xgboost")
             local_xgb_version = version.parse(local_dist.version)
         except importlib_metadata.PackageNotFoundError:
             pass
@@ -138,21 +139,35 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            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
-            )
+            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
+            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output.task)
             if enable_explainability:
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
                     explain_method="explain",
                     target_method="predict",
-                    output_return_type=model_objective_and_output.output_type,
+                    output_return_type=model_task_and_output.output_type,
                 )
                 model_meta.function_properties = {
                     "explain": {model_meta_schema.FunctionProperties.PARTITIONED.value: False}
                 }
 
+                explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
+
+                background_data = handlers_utils.get_explainability_supported_background(
+                    sample_input_data, model_meta, explain_target_method
+                )
+                if background_data is not None:
+                    handlers_utils.save_background_data(
+                        model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, cls.BG_DATA_FILE_SUFFIX, name, background_data
+                    )
+                else:
+                    warnings.warn(
+                        "sample_input_data should be provided for better explainability results",
+                        category=UserWarning,
+                        stacklevel=1,
+                    )
+
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
         model.save_model(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))