snowflake-ml-python 1.6.1__py3-none-any.whl → 1.6.3__py3-none-any.whl

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_meta_schema.ModelObjective:
-        import sklearn.pipeline
-        from sklearn.base import is_classifier, is_regressor
-
-        if isinstance(model, sklearn.pipeline.Pipeline):
-            return model_meta_schema.ModelObjective.UNKNOWN
-        if is_regressor(model):
-            return model_meta_schema.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_meta_schema.ModelObjective.MULTI_CLASSIFICATION
-                return model_meta_schema.ModelObjective.BINARY_CLASSIFICATION
-            return model_meta_schema.ModelObjective.UNKNOWN
-        return model_meta_schema.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(
@@ -106,32 +92,17 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.SKLModelSaveOptions],
     ) -> None:
-        enable_explainability = kwargs.get("enable_explainability", False)
+        # 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)
-
-        enable_explainability = kwargs.get("enable_explainability", False)
         if enable_explainability:
-            # TODO: Currently limited to pandas df, need to extend to other types.
-            if sample_input_data is None or not (
-                isinstance(sample_input_data, pd.DataFrame) or isinstance(sample_input_data, sp_df.DataFrame)
-            ):
-                raise ValueError(
-                    "Sample input data is required to enable explainability. Currently we only support this for "
-                    + "`pandas.DataFrame` and `snowflake.snowpark.dataframe.DataFrame`."
-                )
-            sample_input_data_pandas = (
-                sample_input_data
-                if isinstance(sample_input_data, pd.DataFrame)
-                else snowpark_handler.SnowparkDataFrameHandler.convert_to_df(sample_input_data)
-            )
-            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:
-                sample_input_data_pandas.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(
@@ -141,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)
@@ -159,15 +131,40 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 get_prediction_fn=get_prediction,
             )
 
+            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
-                if cls.get_model_objective(model) == model_meta_schema.ModelObjective.MULTI_CLASSIFICATION:
-                    output_type = model_signature.DataType.STRING
+                handlers_utils.save_background_data(
+                    model_blobs_dir_path,
+                    cls.EXPLAIN_ARTIFACTS_DIR,
+                    cls.BG_DATA_FILE_SUFFIX,
+                    name,
+                    background_data,
+                )
+
                 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)
@@ -184,13 +181,12 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         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")],
-                check_local_version=True,
-            )
+            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
+            [model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn")],
+            check_local_version=True,
         )
 
     @classmethod

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

@@ -1,20 +1,27 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, final
 
 import cloudpickle
 import numpy as np
 import pandas as pd
+from packaging import version
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import type_utils
+from snowflake.ml._internal.exceptions import exceptions
 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
+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, utils as model_signature_utils
 
@@ -36,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
@@ -62,6 +71,60 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
 
         return cast("BaseEstimator", model)
 
+    @classmethod
+    def _get_local_version_package(cls, pkg_name: str) -> Optional[version.Version]:
+        from importlib import metadata as importlib_metadata
+
+        from packaging import version
+
+        local_version = None
+
+        try:
+            local_dist = importlib_metadata.distribution(pkg_name)
+            local_version = version.parse(local_dist.version)
+        except importlib_metadata.PackageNotFoundError:
+            pass
+
+        return local_version
+
+    @classmethod
+    def _can_support_xgb(cls, enable_explainability: Optional[bool]) -> bool:
+
+        local_xgb_version = cls._get_local_version_package("xgboost")
+
+        if local_xgb_version and local_xgb_version >= version.parse("2.1.0"):
+            if enable_explainability:
+                warnings.warn(
+                    f"This version of xgboost {local_xgb_version} does not work with shap 0.42.1."
+                    + "If you want model explanations, lower the xgboost version to <2.1.0.",
+                    category=UserWarning,
+                    stacklevel=1,
+                )
+            return False
+        return True
+
+    @classmethod
+    def _get_supported_object_for_explainability(
+        cls, estimator: "BaseEstimator", enable_explainability: Optional[bool]
+    ) -> Any:
+        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:
+                    result = getattr(estimator, method_name)()
+                    if method_name == "to_xgboost" and not cls._can_support_xgb(enable_explainability):
+                        return None
+                    return result
+                except exceptions.SnowflakeMLException:
+                    pass  # Do nothing and continue to the next method
+        return None
+
     @classmethod
     def save_model(
         cls,
@@ -73,9 +136,8 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.SNOWModelSaveOptions],
     ) -> None:
-        enable_explainability = kwargs.get("enable_explainability", False)
-        if enable_explainability:
-            raise NotImplementedError("Explainability is not supported for Snowpark ML model.")
+
+        enable_explainability = kwargs.get("enable_explainability", None)
 
         from snowflake.ml.modeling.framework.base import BaseEstimator
 
@@ -83,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,
@@ -105,6 +167,35 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                         raise ValueError(f"Target method {method_name} does not exist in the model.")
                 model_meta.signatures = temp_model_signature_dict
 
+        if enable_explainability or enable_explainability is None:
+            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 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_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=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:
@@ -122,7 +213,29 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         model_dependencies = model._get_dependencies()
         for dep in model_dependencies:
             pkg_name = dep.split("==")[0]
-            _include_if_absent_pkgs.append(model_env.ModelDependency(requirement=pkg_name, pip_name=pkg_name))
+            if pkg_name != "xgboost":
+                _include_if_absent_pkgs.append(model_env.ModelDependency(requirement=pkg_name, pip_name=pkg_name))
+                continue
+
+            local_xgb_version = cls._get_local_version_package("xgboost")
+            if local_xgb_version and local_xgb_version >= version.parse("2.0.0") and enable_explainability:
+                model_meta.env.include_if_absent(
+                    [
+                        model_env.ModelDependency(requirement="xgboost==2.0.*", pip_name="xgboost"),
+                    ],
+                    check_local_version=False,
+                )
+            else:
+                model_meta.env.include_if_absent(
+                    [
+                        model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
+                    ],
+                    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
         model_meta.env.include_if_absent(_include_if_absent_pkgs, check_local_version=True)
 
     @classmethod
@@ -163,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:
@@ -177,11 +291,29 @@ class SnowMLModelHandler(_base.BaseModelHandler["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
+
+                    methods = ["to_xgboost", "to_lightgbm", "to_sklearn"]
+                    for method_name in methods:
+                        try:
+                            base_model = getattr(raw_model, method_name)()
+                            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, lightgbm or sklearn (not pipeline) estimator.")
+
+                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)
 
             _SnowMLModel = type(
                 "_SnowMLModel",

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

@@ -111,7 +111,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         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:
-            torch.jit.save(model, f)  # type:ignore[attr-defined]
+            torch.jit.save(model, f)  # type:ignore[no-untyped-call, attr-defined]
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
@@ -141,7 +141,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         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(  # type:ignore[attr-defined]
+            m = torch.jit.load(  # type:ignore[no-untyped-call, attr-defined]
                 f, map_location="cuda" if kwargs.get("use_gpu", False) else "cpu"
             )
         assert isinstance(m, torch.jit.ScriptModule)  # type:ignore[attr-defined]

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

@@ -1,6 +1,7 @@
 # mypy: disable-error-code="import"
-import json
 import os
+import warnings
+from importlib import metadata as importlib_metadata
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -15,12 +16,17 @@ from typing import (
 
 import numpy as np
 import pandas as pd
+from packaging import version
 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,
@@ -47,41 +53,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
 
     MODEL_BLOB_FILE_OR_DIR = "model.ubj"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
-    _BINARY_CLASSIFICATION_OBJECTIVE_PREFIX = ["binary:"]
-    _MULTI_CLASSIFICATION_OBJECTIVE_PREFIX = ["multi:"]
-    _RANKING_OBJECTIVE_PREFIX = ["rank:"]
-    _REGRESSION_OBJECTIVE_PREFIX = ["reg:"]
-
-    @classmethod
-    def get_model_objective(
-        cls, model: Union["xgboost.Booster", "xgboost.XGBModel"]
-    ) -> model_meta_schema.ModelObjective:
-        import xgboost
-
-        if isinstance(model, xgboost.XGBClassifier) or isinstance(model, xgboost.XGBRFClassifier):
-            num_classes = handlers_utils.get_num_classes_if_exists(model)
-            if num_classes == 2:
-                return model_meta_schema.ModelObjective.BINARY_CLASSIFICATION
-            return model_meta_schema.ModelObjective.MULTI_CLASSIFICATION
-        if isinstance(model, xgboost.XGBRegressor) or isinstance(model, xgboost.XGBRFRegressor):
-            return model_meta_schema.ModelObjective.REGRESSION
-        if isinstance(model, xgboost.XGBRanker):
-            return model_meta_schema.ModelObjective.RANKING
-        model_params = json.loads(model.save_config())
-        model_objective = model_params["learner"]["objective"]
-        for classification_objective in cls._BINARY_CLASSIFICATION_OBJECTIVE_PREFIX:
-            if classification_objective in model_objective:
-                return model_meta_schema.ModelObjective.BINARY_CLASSIFICATION
-        for classification_objective in cls._MULTI_CLASSIFICATION_OBJECTIVE_PREFIX:
-            if classification_objective in model_objective:
-                return model_meta_schema.ModelObjective.MULTI_CLASSIFICATION
-        for ranking_objective in cls._RANKING_OBJECTIVE_PREFIX:
-            if ranking_objective in model_objective:
-                return model_meta_schema.ModelObjective.RANKING
-        for regression_objective in cls._REGRESSION_OBJECTIVE_PREFIX:
-            if regression_objective in model_objective:
-                return model_meta_schema.ModelObjective.REGRESSION
-        return model_meta_schema.ModelObjective.UNKNOWN
+    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba"]
 
     @classmethod
     def can_handle(
@@ -116,10 +88,29 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.XGBModelSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", True)
+
         import xgboost
 
         assert isinstance(model, xgboost.Booster) or isinstance(model, xgboost.XGBModel)
 
+        local_xgb_version = None
+
+        try:
+            local_dist = importlib_metadata.distribution("xgboost")
+            local_xgb_version = version.parse(local_dist.version)
+        except importlib_metadata.PackageNotFoundError:
+            pass
+
+        if local_xgb_version and local_xgb_version >= version.parse("2.1.0") and enable_explainability:
+            warnings.warn(
+                f"This version of xgboost {local_xgb_version} does not work with shap 0.42.1."
+                + "If you want model explanations, lower the xgboost version to <2.1.0.",
+                category=UserWarning,
+                stacklevel=1,
+            )
+            enable_explainability = False
+
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
                 model=model,
@@ -148,22 +139,35 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_objective = cls.get_model_objective(model)
-            model_meta.model_objective = model_objective
-            if kwargs.get("enable_explainability", True):
-                output_type = model_signature.DataType.DOUBLE
-                if model_objective == model_meta_schema.ModelObjective.MULTI_CLASSIFICATION:
-                    output_type = model_signature.DataType.STRING
+            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=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))
@@ -180,15 +184,26 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         model_meta.env.include_if_absent(
             [
                 model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn"),
-                model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
             ],
             check_local_version=True,
         )
-        if kwargs.get("enable_explainability", True):
+        if local_xgb_version and local_xgb_version >= version.parse("2.0.0") and enable_explainability:
+            model_meta.env.include_if_absent(
+                [
+                    model_env.ModelDependency(requirement="xgboost==2.0.*", pip_name="xgboost"),
+                ],
+                check_local_version=False,
+            )
+        else:
             model_meta.env.include_if_absent(
-                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
+                [
+                    model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
+                ],
                 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
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
@@ -269,7 +284,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     import shap
 
                     explainer = shap.TreeExplainer(raw_model)
-                    df = pd.DataFrame(explainer(X).values)
+                    df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
                 if target_method == "explain":