snowflake-ml-python 1.7.2__py3-none-any.whl → 1.7.3__py3-none-any.whl

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

@@ -195,8 +195,12 @@ class HuggingFacePipelineHandler(
         os.makedirs(model_blob_path, exist_ok=True)
 
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
+            save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             model.save_pretrained(  # type:ignore[attr-defined]
-                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+                save_path
+            )
+            handlers_utils.save_transformers_config_with_auto_map(
+                save_path,
             )
             pipeline_params = {
                 "_batch_size": model._batch_size,  # type:ignore[attr-defined]

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

@@ -110,8 +110,8 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_task_and_output = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output.task)
+            model_task_and_output = model_task_utils.resolve_model_task_and_output_type(model, model_meta.task)
+            model_meta.task = model_task_and_output.task
             if enable_explainability:
                 explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(
@@ -240,7 +240,9 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                     import shap
 
                     explainer = shap.TreeExplainer(raw_model)
-                    df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
+                    df = handlers_utils.convert_explanations_to_2D_df(
+                        raw_model, explainer.shap_values(X, from_call=True)
+                    )
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
                 if target_method == "explain":

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

@@ -14,8 +14,8 @@ 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 utils as model_signature_utils
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 if TYPE_CHECKING:
@@ -24,6 +24,25 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _validate_sentence_transformers_signatures(sigs: Dict[str, model_signature.ModelSignature]) -> None:
+    if list(sigs.keys()) != ["encode"]:
+        raise ValueError("target_methods can only be ['encode']")
+
+    if len(sigs["encode"].inputs) != 1:
+        raise ValueError("SentenceTransformer can only accept 1 input column")
+
+    if len(sigs["encode"].outputs) != 1:
+        raise ValueError("SentenceTransformer can only return 1 output column")
+
+    assert isinstance(sigs["encode"].inputs[0], model_signature.FeatureSpec)
+
+    if sigs["encode"].inputs[0]._shape is not None:
+        raise ValueError("SentenceTransformer does not support input shape")
+
+    if sigs["encode"].inputs[0]._dtype != model_signature.DataType.STRING:
+        raise ValueError("SentenceTransformer only accepts string input")
+
+
 @final
 class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.SentenceTransformer"]):
     HANDLER_TYPE = "sentence_transformers"
@@ -68,6 +87,10 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         if enable_explainability:
             raise NotImplementedError("Explainability is not supported for Sentence Transformer model.")
 
+        batch_size = kwargs.get("batch_size", 32)
+        if not isinstance(batch_size, int) or batch_size <= 0:
+            raise ValueError("batch_size must be a positive integer")
+
         # Validate target methods and signature (if possible)
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -75,12 +98,23 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
                 target_methods=kwargs.pop("target_methods", None),
                 default_target_methods=cls.DEFAULT_TARGET_METHODS,
             )
-            assert target_methods == ["encode"], "target_methods can only be ['encode']"
+            if target_methods != ["encode"]:
+                raise ValueError("target_methods can only be ['encode']")
 
             def get_prediction(
                 target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
             ) -> model_types.SupportedLocalDataType:
-                return _sentence_transformer_encode(model, sample_input_data)
+                if not isinstance(sample_input_data, pd.DataFrame):
+                    sample_input_data = model_signature._convert_local_data_to_df(data=sample_input_data)
+
+                if sample_input_data.shape[1] != 1:
+                    raise ValueError(
+                        "SentenceTransformer can only accept 1 input column when converted to pd.DataFrame"
+                    )
+                X_list = sample_input_data.iloc[:, 0].tolist()
+
+                assert callable(getattr(model, "encode", None))
+                return pd.DataFrame({0: model.encode(X_list, batch_size=batch_size).tolist()})
 
             if model_meta.signatures:
                 handlers_utils.validate_target_methods(model, list(model_meta.signatures.keys()))
@@ -102,10 +136,16 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
                         get_prediction_fn=get_prediction,
                     )
 
+            _validate_sentence_transformers_signatures(model_meta.signatures)
+
         # 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.MODEL_BLOB_FILE_OR_DIR))
+        save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+        model.save(save_path)
+        handlers_utils.save_transformers_config_with_auto_map(
+            save_path,
+        )
 
         # save model metadata
         base_meta = model_blob_meta.ModelBlobMeta(
@@ -113,6 +153,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
+            options=model_meta_schema.SentenceTransformersModelBlobOptions(batch_size=batch_size),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -149,6 +190,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
             # We need to redirect the same folders to a writable location in the sandbox.
             os.environ["TRANSFORMERS_CACHE"] = "/tmp"
+            os.environ["HF_HOME"] = "/tmp"
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         model_blobs_metadata = model_meta.models
@@ -183,6 +225,10 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
             raw_model: "sentence_transformers.SentenceTransformer",
             model_meta: model_meta_api.ModelMetadata,
         ) -> Type[custom_model.CustomModel]:
+            batch_size = cast(
+                model_meta_schema.SentenceTransformersModelBlobOptions, model_meta.models[model_meta.name].options
+            ).get("batch_size", None)
+
             def get_prediction(
                 raw_model: "sentence_transformers.SentenceTransformer",
                 signature: model_signature.ModelSignature,
@@ -190,8 +236,11 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
-                    predictions_df = _sentence_transformer_encode(raw_model, X)
-                    return model_signature_utils.rename_pandas_df(predictions_df, signature.outputs)
+                    X_list = X.iloc[:, 0].tolist()
+
+                    return pd.DataFrame(
+                        {signature.outputs[0].name: raw_model.encode(X_list, batch_size=batch_size).tolist()}
+                    )
 
                 return fn
 
@@ -217,17 +266,3 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         predict_method = getattr(sentence_transformers_SentenceTransformer_model, "encode", None)
         assert callable(predict_method)
         return sentence_transformers_SentenceTransformer_model
-
-
-def _sentence_transformer_encode(
-    model: "sentence_transformers.SentenceTransformer", X: model_types.SupportedLocalDataType
-) -> model_types.SupportedLocalDataType:
-
-    if not isinstance(X, pd.DataFrame):
-        X = model_signature._convert_local_data_to_df(X)
-
-    assert X.shape[1] == 1, "SentenceTransformer can only accept 1 input column when converted to pd.DataFrame"
-    X_list = X.iloc[:, 0].tolist()
-
-    assert callable(getattr(model, "encode", None))
-    return pd.DataFrame({0: model.encode(X_list, batch_size=X.shape[0]).tolist()})

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

@@ -152,8 +152,8 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 sample_input_data, model_meta, explain_target_method
             )
 
-            model_task_and_output_type = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output_type.task)
+            model_task_and_output_type = model_task_utils.resolve_model_task_and_output_type(model, model_meta.task)
+            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:
@@ -164,11 +164,17 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                         stacklevel=1,
                     )
                     enable_explainability = False
-                elif model_meta.task == model_types.Task.UNKNOWN:
+                elif model_meta.task == model_types.Task.UNKNOWN or explain_target_method is None:
                     enable_explainability = False
                 else:
                     enable_explainability = True
             if enable_explainability:
+                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,
+                )
                 handlers_utils.save_background_data(
                     model_blobs_dir_path,
                     cls.EXPLAIN_ARTIFACTS_DIR,
@@ -177,13 +183,6 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                     background_data,
                 )
 
-                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,
-                )
-
         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:

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

@@ -68,21 +68,45 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         return cast("BaseEstimator", model)
 
     @classmethod
-    def _get_supported_object_for_explainability(cls, estimator: "BaseEstimator") -> Any:
+    def _get_supported_object_for_explainability(
+        cls,
+        estimator: "BaseEstimator",
+        background_data: Optional[model_types.SupportedDataType],
+        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:
+        tree_methods = ["to_xgboost", "to_lightgbm", "to_sklearn"]
+        non_tree_methods = ["to_sklearn"]
+        for method_name in tree_methods:
+            if hasattr(estimator, method_name):
+                try:
+                    result = getattr(estimator, method_name)()
+                    return result
+                except exceptions.SnowflakeMLException:
+                    pass  # Do nothing and continue to the next method
+        for method_name in non_tree_methods:
             if hasattr(estimator, method_name):
                 try:
                     result = getattr(estimator, method_name)()
+                    if enable_explainability is None and background_data is None:
+                        return None  # cannot get explain without background data
+                    elif enable_explainability and background_data is None:
+                        raise ValueError(
+                            "Provide `sample_input_data` to generate explanations for sklearn Snowpark ML models."
+                        )
                     return result
                 except exceptions.SnowflakeMLException:
                     pass  # Do nothing and continue to the next method
+
+        if enable_explainability:
+            raise ValueError(
+                "Explain only supported for xgboost, lightgbm and sklearn (not pipeline) Snowpark ML models."
+            )
         return None
 
     @classmethod
@@ -127,34 +151,39 @@ 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)
-            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."
-                    )
+        python_base_obj = cls._get_supported_object_for_explainability(model, sample_input_data, enable_explainability)
+        explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
+
+        if enable_explainability:
+            if explain_target_method is None:
+                raise ValueError(
+                    "The model must have one of the following methods to enable explainability: "
+                    + ", ".join(cls.EXPLAIN_TARGET_METHODS)
+                )
+        if enable_explainability is None:
+            if python_base_obj is None or explain_target_method is None:
                 # set None to False so we don't include shap in the environment
                 enable_explainability = False
             else:
-                model_task_and_output_type = model_task_utils.get_model_task_and_output_type(python_base_obj)
-                model_meta.task = handlers_utils.validate_model_task(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 enable_explainability:
+            model_task_and_output_type = model_task_utils.resolve_model_task_and_output_type(
+                python_base_obj, model_meta.task
+            )
+            model_meta.task = model_task_and_output_type.task
+            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,
+            )
+            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
                 )
-                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)
@@ -237,8 +266,17 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 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:
+                    tree_methods = ["to_xgboost", "to_lightgbm"]
+                    non_tree_methods = ["to_sklearn"]
+                    for method_name in tree_methods:
+                        try:
+                            base_model = getattr(raw_model, method_name)()
+                            explainer = shap.TreeExplainer(base_model)
+                            df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer.shap_values(X))
+                            return model_signature_utils.rename_pandas_df(df, signature.outputs)
+                        except exceptions.SnowflakeMLException:
+                            pass  # Do nothing and continue to the next method
+                    for method_name in non_tree_methods:
                         try:
                             base_model = getattr(raw_model, method_name)()
                             explainer = shap.Explainer(base_model, masker=background_data)

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

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
 
 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
@@ -73,13 +74,42 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         if enable_explainability:
             raise NotImplementedError("Explainability is not supported for Tensorflow model.")
 
+        # When tensorflow is installed, keras is also installed.
+        import keras
         import tensorflow
 
         assert isinstance(model, tensorflow.Module)
 
         is_keras_model = type_utils.LazyType("tensorflow.keras.Model").isinstance(model) or type_utils.LazyType(
-            "tf_keras.Model"
+            "keras.Model"
         ).isinstance(model)
+        is_tf_keras_model = type_utils.LazyType("tf_keras.Model").isinstance(model)
+        is_keras_functional_or_sequential_model = (
+            getattr(model, "_is_graph_network", False)
+            or type_utils.LazyType("tensorflow.keras.engine.sequential.Sequential").isinstance(model)
+            or type_utils.LazyType("keras.engine.sequential.Sequential").isinstance(model)
+            or type_utils.LazyType("tf_keras.engine.sequential.Sequential").isinstance(model)
+        )
+
+        assert isinstance(model, tensorflow.Module)
+
+        keras_version = version.parse(keras.__version__)
+
+        # Tensorflow and keras model save format is different.
+        # Keras functional or sequential models are saved as keras format
+        # Keras v3 other models are saved using cloudpickle
+        # Keras v2 other models are saved using tensorflow saved model format
+        # Tensorflow models are saved using tensorflow saved model format
+
+        if is_keras_model or is_tf_keras_model:
+            if is_keras_functional_or_sequential_model:
+                save_format = "keras"
+            elif keras_version.major == 2 or is_tf_keras_model:
+                save_format = "keras_tf"
+            else:
+                save_format = "cloudpickle"
+        else:
+            save_format = "tf"
 
         if is_keras_model:
             default_target_methods = ["predict"]
@@ -93,6 +123,9 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
                 default_target_methods=default_target_methods,
             )
 
+            if is_keras_model and len(target_methods) > 1:
+                raise ValueError("Keras model can only have one target method.")
+
             def get_prediction(
                 target_method_name: str, sample_input_data: "model_types.SupportedLocalDataType"
             ) -> model_types.SupportedLocalDataType:
@@ -122,31 +155,43 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
-        if is_keras_model:
-            tensorflow.keras.models.save_model(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
-            model_meta.env.include_if_absent(
-                [
-                    model_env.ModelDependency(requirement="keras<3", pip_name="keras"),
-                ],
-                check_local_version=False,
-            )
+        save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+        if save_format == "keras":
+            model.save(save_path, save_format="keras")
+        elif save_format == "keras_tf":
+            model.save(save_path, save_format="tf")
+        elif save_format == "cloudpickle":
+            import cloudpickle
+
+            with open(save_path, "wb") as f:
+                cloudpickle.dump(model, f)
         else:
-            tensorflow.saved_model.save(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
+            tensorflow.saved_model.save(
+                model,
+                save_path,
+                options=tensorflow.saved_model.SaveOptions(experimental_custom_gradients=False),
+            )
 
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
-            options=model_meta_schema.TensorflowModelBlobOptions(is_keras_model=is_keras_model),
+            options=model_meta_schema.TensorflowModelBlobOptions(save_format=save_format),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
+        dependencies = [
+            model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"),
+        ]
+        if is_keras_model:
+            dependencies.append(model_env.ModelDependency(requirement="keras", pip_name="keras"))
+        elif is_tf_keras_model:
+            dependencies.append(model_env.ModelDependency(requirement="tf-keras", pip_name="tf-keras"))
+
         model_meta.env.include_if_absent(
-            [
-                model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"),
-            ],
+            dependencies,
             check_local_version=True,
         )
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
@@ -166,10 +211,18 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blob_metadata = model_blobs_metadata[name]
         model_blob_filename = model_blob_metadata.path
         model_blob_options = cast(model_meta_schema.TensorflowModelBlobOptions, model_blob_metadata.options)
-        if model_blob_options.get("is_keras_model", False):
-            m = tensorflow.keras.models.load_model(os.path.join(model_blob_path, model_blob_filename), compile=False)
+        load_path = os.path.join(model_blob_path, model_blob_filename)
+        save_format = model_blob_options.get("save_format", "tf")
+        if save_format == "keras" or save_format == "keras_tf":
+            m = tensorflow.keras.models.load_model(load_path)
+        elif save_format == "cloudpickle":
+            import cloudpickle
+
+            with open(load_path, "rb") as f:
+                m = cloudpickle.load(f)
         else:
-            m = tensorflow.saved_model.load(os.path.join(model_blob_path, model_blob_filename))
+            m = tensorflow.saved_model.load(load_path)
+
         return cast(tensorflow.Module, m)
 
     @classmethod

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

@@ -117,8 +117,8 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_task_and_output = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output.task)
+            model_task_and_output = model_task_utils.resolve_model_task_and_output_type(model, model_meta.task)
+            model_meta.task = model_task_and_output.task
             if enable_explainability:
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
@@ -254,7 +254,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     import shap
 
                     explainer = shap.TreeExplainer(raw_model)
-                    df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
+                    df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer.shap_values(X))
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
                 if target_method == "explain":

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -215,6 +215,7 @@ class ModelMetadata:
         function_properties: A dict mapping function names to dict mapping function property key to value.
         metadata: User provided key-value metadata of the model. Defaults to None.
         creation_timestamp: Unix timestamp when the model metadata is created.
+        user_files: Dict mapping subdirectories to extra artifact file paths for files to include in the model.
         task: Model task like TABULAR_REGRESSION, tabular_classification, timeseries_forecasting etc.
     """
 
@@ -234,6 +235,7 @@ class ModelMetadata:
         runtimes: Optional[Dict[str, model_runtime.ModelRuntime]] = None,
         signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
         function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+        user_files: Optional[Dict[str, List[str]]] = None,
         metadata: Optional[Dict[str, str]] = None,
         creation_timestamp: Optional[str] = None,
         min_snowpark_ml_version: Optional[str] = None,
@@ -247,6 +249,7 @@ class ModelMetadata:
         if signatures:
             self.signatures = signatures
         self.function_properties = function_properties or {}
+        self.user_files = user_files
         self.metadata = metadata
         self.model_type = model_type
         self.env = env

snowflake/ml/model/_packager/model_meta/model_meta_schema.py

@@ -59,7 +59,11 @@ class XgboostModelBlobOptions(BaseModelBlobOptions):
 
 
 class TensorflowModelBlobOptions(BaseModelBlobOptions):
-    is_keras_model: Required[bool]
+    save_format: Required[str]
+
+
+class SentenceTransformersModelBlobOptions(BaseModelBlobOptions):
+    batch_size: Required[int]
 
 
 ModelBlobOptions = Union[
@@ -68,6 +72,7 @@ ModelBlobOptions = Union[
     MLFlowModelBlobOptions,
     XgboostModelBlobOptions,
     TensorflowModelBlobOptions,
+    SentenceTransformersModelBlobOptions,
 ]
 
 

snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py

@@ -1,2 +1,2 @@
-REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2022.11,<2024', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2022.11,<2024', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
-ALL_REQUIREMENTS=['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2022.11,<2024', 'importlib_resources>=6.1.1, <7', 'lightgbm>=4.1.0, <5', 'mlflow>=2.1.0,<2.4', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pytorch>=2.0.1,<2.3.0', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2022.11,<2024', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'sentence-transformers>=2.2.2,<3', 'sentencepiece>=0.1.95,<1', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'tensorflow>=2.10,<3', 'tokenizers>=0.10,<1', 'torchdata>=0.4,<1', 'transformers>=4.32.1,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
+REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
+ALL_REQUIREMENTS=['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'huggingface_hub<0.26', 'importlib_resources>=6.1.1, <7', 'lightgbm>=4.1.0, <5', 'mlflow>=2.16.0, <3', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pytorch>=2.0.1,<2.3.0', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'sentence-transformers>=2.2.2,<3', 'sentencepiece>=0.1.95,<1', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'tensorflow>=2.12.0,<3', 'tokenizers>=0.10,<1', 'torchdata>=0.4,<1', 'transformers>=4.32.1,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']

snowflake/ml/model/_packager/model_task/model_task_utils.py

@@ -149,8 +149,9 @@ def _get_model_task(model: Any) -> type_hints.Task:
     raise ValueError(f"Model type {type(model)} is not supported")
 
 
-def get_model_task_and_output_type(model: Any) -> ModelTaskAndOutputType:
-    task = _get_model_task(model)
+def resolve_model_task_and_output_type(model: Any, passed_model_task: type_hints.Task) -> ModelTaskAndOutputType:
+    inferred_task = _get_model_task(model)
+    task = handlers_utils.validate_model_task(passed_model_task, inferred_task)
     output_type = model_signature.DataType.DOUBLE
     if task == type_hints.Task.TABULAR_MULTI_CLASSIFICATION:
         output_type = model_signature.DataType.STRING

snowflake/ml/model/_signatures/pandas_handler.py

@@ -224,6 +224,6 @@ class PandasDataFrameHandler(base_handler.BaseDataHandler[pd.DataFrame]):
         df_col_dtypes = [df[col].dtype for col in df.columns]
         for df_col, df_col_dtype in zip(df_cols, df_col_dtypes):
             if df_col_dtype == np.dtype("O"):
-                if isinstance(df[df_col][0], np.ndarray):
+                if isinstance(df[df_col].iloc[0], np.ndarray):
                     df[df_col] = df[df_col].map(np.ndarray.tolist)
         return df

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -1,5 +1,5 @@
 import json
-from typing import Literal, Optional, Sequence, cast
+from typing import Any, Literal, Optional, Sequence, cast
 
 import numpy as np
 import pandas as pd
@@ -73,14 +73,20 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
                 assert isinstance(feature, core.FeatureSpec), "Invalid feature kind."
                 dtype_map[feature.name] = feature.as_dtype()
         df_local = data.to_pandas()
+
         # This is because Array will become string (Even though the correct schema is set)
         # and object will become variant type and requires an additional loads
         # to get correct data otherwise it would be string.
+        def load_if_not_null(x: str) -> Optional[Any]:
+            if x is None:
+                return None
+            return json.loads(x)
+
         for field in data.schema.fields:
             if isinstance(field.datatype, spt.ArrayType):
                 df_local[identifier.get_unescaped_names(field.name)] = df_local[
                     identifier.get_unescaped_names(field.name)
-                ].map(json.loads)
+                ].map(load_if_not_null)
         # Only when the feature is not from inference, we are confident to do the type casting.
         # Otherwise, dtype_map will be empty.
         # Errors are ignored to make sure None won't be converted and won't raise Error

snowflake/ml/model/type_hints.py

@@ -199,6 +199,7 @@ class HuggingFaceSaveOptions(BaseModelSaveOption):
 class SentenceTransformersSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    batch_size: NotRequired[int]
 
 
 ModelSaveOption = Union[