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

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

@@ -1,6 +1,8 @@
+import os
 from abc import abstractmethod
 from typing import Dict, Generic, Optional, Protocol, Type, final
 
+import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml.model import custom_model, type_hints as model_types
@@ -16,7 +18,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
 
     @classmethod
     @abstractmethod
-    def can_handle(cls, model: model_types.SupportedDataType) -> TypeGuard[model_types._ModelType]:
+    def can_handle(cls, model: model_types.SupportedModelType) -> TypeGuard[model_types._ModelType]:
         """Whether this handler could support the type of the `model`.
 
         Args:
@@ -96,6 +98,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
         cls,
         raw_model: model_types._ModelType,
         model_meta: model_meta.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.BaseModelLoadOption],
     ) -> custom_model.CustomModel:
         """Create a custom model class wrap for unified interface when being deployed. The predict method will be
@@ -104,6 +107,7 @@ class _BaseModelHandlerProtocol(Protocol[model_types._ModelType]):
         Args:
             raw_model: original model object,
             model_meta: The model metadata.
+            background_data: The background data used for the model explanations.
             kwargs: Options when converting the model.
 
         Raises:
@@ -121,7 +125,8 @@ class BaseModelHandler(Generic[model_types._ModelType], _BaseModelHandlerProtoco
     _MIN_SNOWPARK_ML_VERSION: The minimal version of Snowpark ML library to use the current handler.
     _HANDLER_MIGRATOR_PLANS: Dict holding handler migrator plans.
 
-    MODELE_BLOB_FILE_OR_DIR: Relative path of the model blob file in the model subdir. Default to "model.pkl".
+    MODEL_BLOB_FILE_OR_DIR: Relative path of the model blob file in the model subdir. Default to "model.pkl".
+    BG_DATA_FILE_SUFFIX: Suffix of the background data file. Default to "_background_data.pqt".
     MODEL_ARTIFACTS_DIR: Relative path of the model artifacts dir in the model subdir. Default to "artifacts"
     DEFAULT_TARGET_METHODS: Default target methods to be logged if not specified in this kind of model. Default to
         ["predict"]
@@ -129,8 +134,10 @@ class BaseModelHandler(Generic[model_types._ModelType], _BaseModelHandlerProtoco
         inputting sample data or model signature. Default to False.
     """
 
-    MODELE_BLOB_FILE_OR_DIR = "model.pkl"
+    MODEL_BLOB_FILE_OR_DIR = "model.pkl"
+    BG_DATA_FILE_SUFFIX = "_background_data.pqt"
     MODEL_ARTIFACTS_DIR = "artifacts"
+    EXPLAIN_ARTIFACTS_DIR = "explain_artifacts"
     DEFAULT_TARGET_METHODS = ["predict"]
     IS_AUTO_SIGNATURE = False
 
@@ -159,3 +166,23 @@ class BaseModelHandler(Generic[model_types._ModelType], _BaseModelHandlerProtoco
                 model_meta=model_meta,
                 model_blobs_dir_path=model_blobs_dir_path,
             )
+
+    @classmethod
+    @final
+    def load_background_data(cls, name: str, model_blobs_dir_path: str) -> Optional[pd.DataFrame]:
+        """Load the model into memory.
+
+        Args:
+            name: Name of the model.
+            model_blobs_dir_path: Directory path to the whole model.
+
+        Returns:
+            Optional[pd.DataFrame], background data as pandas DataFrame, if exists.
+        """
+        data_blob_path = os.path.join(model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, name + cls.BG_DATA_FILE_SUFFIX)
+        if not os.path.exists(model_blobs_dir_path) or not os.path.isfile(data_blob_path):
+            return None
+        with open(data_blob_path, "rb") as f:
+            background_data = pd.read_parquet(f)
+
+        return background_data

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

@@ -1,4 +1,9 @@
-from typing import Callable, Iterable, Optional, Sequence, cast
+import json
+from typing import Any, Callable, Iterable, Optional, Sequence, cast
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
 
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_meta import model_meta
@@ -40,6 +45,24 @@ def validate_signature(
     return model_meta
 
 
+def add_explain_method_signature(
+    model_meta: model_meta.ModelMetadata,
+    explain_method: str,
+    target_method: str,
+    output_return_type: model_signature.DataType = model_signature.DataType.DOUBLE,
+) -> model_meta.ModelMetadata:
+    if target_method not in model_meta.signatures:
+        raise ValueError(f"Signature for target method {target_method} is missing")
+    inputs = model_meta.signatures[target_method].inputs
+    model_meta.signatures[explain_method] = model_signature.ModelSignature(
+        inputs=inputs,
+        outputs=[
+            model_signature.FeatureSpec(dtype=output_return_type, name=f"{spec.name}_explanation") for spec in inputs
+        ],
+    )
+    return model_meta
+
+
 def get_target_methods(
     model: model_types.SupportedModelType,
     target_methods: Optional[Sequence[str]],
@@ -56,3 +79,37 @@ def validate_target_methods(model: model_types.SupportedModelType, target_method
     for method_name in target_methods:
         if not _is_callable(model, method_name):
             raise ValueError(f"Target method {method_name} is not callable or does not exist in the model.")
+
+
+def get_num_classes_if_exists(model: model_types.SupportedModelType) -> int:
+    num_classes = getattr(model, "classes_", [])
+    return len(num_classes)
+
+
+def convert_explanations_to_2D_df(
+    model: model_types.SupportedModelType, explanations: npt.NDArray[Any]
+) -> pd.DataFrame:
+    if explanations.ndim != 3:
+        return pd.DataFrame(explanations)
+
+    if hasattr(model, "classes_"):
+        classes_list = [cl for cl in model.classes_]  # type:ignore[union-attr]
+        len_classes = len(classes_list)
+        if explanations.shape[2] != len_classes:
+            raise ValueError(f"Model has {len_classes} classes but explanations have {explanations.shape[2]}")
+    else:
+        classes_list = [i for i in range(explanations.shape[2])]
+    exp_2d = []
+    # TODO (SNOW-1549044): Optimize this
+    for row in explanations:
+        col_list = []
+        for column in row:
+            class_explanations = {}
+            for cl, cl_exp in zip(classes_list, column):
+                if isinstance(cl, (int, np.integer)):
+                    cl = int(cl)
+                class_explanations[cl] = cl_exp
+            col_list.append(json.dumps(class_explanations))
+        exp_2d.append(col_list)
+
+    return pd.DataFrame(exp_2d)

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

@@ -30,9 +30,25 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
-    MODELE_BLOB_FILE_OR_DIR = "model.bin"
+    MODEL_BLOB_FILE_OR_DIR = "model.bin"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
 
+    @classmethod
+    def get_model_objective(cls, model: "catboost.CatBoost") -> model_meta_schema.ModelObjective:
+        import catboost
+
+        if isinstance(model, catboost.CatBoostClassifier):
+            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, catboost.CatBoostRanker):
+            return model_meta_schema.ModelObjective.RANKING
+        if isinstance(model, catboost.CatBoostRegressor):
+            return model_meta_schema.ModelObjective.REGRESSION
+        # TODO: Find out model type from the generic Catboost Model
+        return model_meta_schema.ModelObjective.UNKNOWN
+
     @classmethod
     def can_handle(cls, model: model_types.SupportedModelType) -> TypeGuard["catboost.CatBoost"]:
         return (type_utils.LazyType("catboost.CatBoost").isinstance(model)) and any(
@@ -89,10 +105,25 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 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_meta = handlers_utils.add_explain_method_signature(
+                    model_meta=model_meta,
+                    explain_method="explain",
+                    target_method="predict",
+                    output_return_type=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)
 
         model.save_model(model_save_path)
 
@@ -100,7 +131,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
             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.CatBoostModelBlobOptions({"catboost_estimator_type": model.__class__.__name__}),
         )
         model_meta.models[name] = base_meta
@@ -112,6 +143,12 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
             ],
             check_local_version=True,
         )
+        if kwargs.get("enable_explainability", True):
+            model_meta.env.include_if_absent(
+                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
+                check_local_version=True,
+            )
+            model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
         return None
@@ -157,6 +194,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
         cls,
         raw_model: "catboost.CatBoost",
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.CatBoostModelLoadOptions],
     ) -> custom_model.CustomModel:
         import catboost
@@ -186,6 +224,17 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
 
                     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
+
+                    explainer = shap.TreeExplainer(raw_model)
+                    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":
+                    return explain_fn
+
                 return fn
 
             type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}

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

@@ -51,6 +51,9 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         **kwargs: Unpack[model_types.CustomModelSaveOption],
     ) -> None:
         assert isinstance(model, custom_model.CustomModel)
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for custom model.")
 
         def get_prediction(
             target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
@@ -108,13 +111,13 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         # Make sure that the module where the model is defined get pickled by value as well.
         cloudpickle.register_pickle_by_value(sys.modules[model.__module__])
         pickled_obj = (model.__class__, model.context)
-        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(pickled_obj, f)
         # model meta will be saved by the context manager
         model_meta.models[name] = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
-            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            path=cls.MODEL_BLOB_FILE_OR_DIR,
             handler_version=cls.HANDLER_VERSION,
             function_properties=model_meta.function_properties,
             artifacts={
@@ -183,6 +186,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
         cls,
         raw_model: custom_model.CustomModel,
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.CustomModelLoadOption],
     ) -> custom_model.CustomModel:
         return raw_model

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

@@ -89,7 +89,7 @@ class HuggingFacePipelineHandler(
     _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"
     ADDITIONAL_CONFIG_FILE = "pipeline_config.pt"
     DEFAULT_TARGET_METHODS = ["__call__"]
     IS_AUTO_SIGNATURE = True
@@ -133,6 +133,9 @@ class HuggingFacePipelineHandler(
         is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.HuggingFaceSaveOptions],
     ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for huggingface model.")
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             task = model.task  # type:ignore[attr-defined]
             framework = model.framework  # type:ignore[attr-defined]
@@ -193,7 +196,7 @@ class HuggingFacePipelineHandler(
 
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             model.save_pretrained(  # type:ignore[attr-defined]
-                os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR)
+                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             )
             pipeline_params = {
                 "_batch_size": model._batch_size,  # type:ignore[attr-defined]
@@ -205,7 +208,7 @@ class HuggingFacePipelineHandler(
             with open(
                 os.path.join(
                     model_blob_path,
-                    cls.MODELE_BLOB_FILE_OR_DIR,
+                    cls.MODEL_BLOB_FILE_OR_DIR,
                     cls.ADDITIONAL_CONFIG_FILE,
                 ),
                 "wb",
@@ -213,7 +216,7 @@ class HuggingFacePipelineHandler(
                 cloudpickle.dump(pipeline_params, f)
         else:
             with open(
-                os.path.join(model_blob_path, cls.MODELE_BLOB_FILE_OR_DIR),
+                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR),
                 "wb",
             ) as f:
                 cloudpickle.dump(model, f)
@@ -222,7 +225,7 @@ class HuggingFacePipelineHandler(
             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.HuggingFacePipelineModelBlobOptions(
                 {
                     "task": task,
@@ -329,6 +332,7 @@ class HuggingFacePipelineHandler(
         cls,
         raw_model: Union[huggingface_pipeline.HuggingFacePipelineModel, "transformers.Pipeline"],
         model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
     ) -> custom_model.CustomModel:
         import transformers

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

@@ -41,8 +41,49 @@ 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"]
+    ) -> model_meta_schema.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 model_meta_schema.ModelObjective.BINARY_CLASSIFICATION
+            return model_meta_schema.ModelObjective.MULTI_CLASSIFICATION
+        if isinstance(model, lightgbm.LGBMRanker):
+            return model_meta_schema.ModelObjective.RANKING
+        if isinstance(model, lightgbm.LGBMRegressor):
+            return model_meta_schema.ModelObjective.REGRESSION
+        model_objective = model.params["objective"]
+        if model_objective in cls._BINARY_CLASSIFICATION_OBJECTIVES:
+            return model_meta_schema.ModelObjective.BINARY_CLASSIFICATION
+        if model_objective in cls._MULTI_CLASSIFICATION_OBJECTIVES:
+            return model_meta_schema.ModelObjective.MULTI_CLASSIFICATION
+        if model_objective in cls._RANKING_OBJECTIVES:
+            return model_meta_schema.ModelObjective.RANKING
+        if model_objective in cls._REGRESSION_OBJECTIVES:
+            return model_meta_schema.ModelObjective.REGRESSION
+        return model_meta_schema.ModelObjective.UNKNOWN
 
     @classmethod
     def can_handle(
@@ -105,11 +146,29 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 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 in [
+                    model_meta_schema.ModelObjective.BINARY_CLASSIFICATION,
+                    model_meta_schema.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_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)
 
@@ -117,7 +176,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
@@ -130,6 +189,12 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
             ],
             check_local_version=True,
         )
+        if kwargs.get("enable_explainability", True):
+            model_meta.env.include_if_absent(
+                [model_env.ModelDependency(requirement="shap", pip_name="shap")],
+                check_local_version=True,
+            )
+            model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
 
         return None
 
@@ -169,6 +234,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
@@ -198,6 +264,17 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
 
                     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
+
+                    explainer = shap.TreeExplainer(raw_model)
+                    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":
+                    return explain_fn
+
                 return fn
 
             type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}

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

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/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