snowflake-ml-python 1.7.4__py3-none-any.whl → 1.7.5__py3-none-any.whl

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

@@ -0,0 +1,230 @@
+import os
+from typing import TYPE_CHECKING, 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.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_migrator import base_migrator
+from snowflake.ml.model._packager.model_meta import (
+    model_blob_meta,
+    model_meta as model_meta_api,
+)
+from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
+
+if TYPE_CHECKING:
+    import keras
+
+
+@final
+class KerasHandler(_base.BaseModelHandler["keras.Model"]):
+    """Handler for Keras v3 model.
+
+    Currently keras.Model based classes are supported.
+    """
+
+    HANDLER_TYPE = "keras"
+    HANDLER_VERSION = "2025-01-01"
+    _MIN_SNOWPARK_ML_VERSION = "1.7.5"
+    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+
+    MODEL_BLOB_FILE_OR_DIR = "model.keras"
+    CUSTOM_OBJECT_SAVE_PATH = "custom_objects.pkl"
+    DEFAULT_TARGET_METHODS = ["predict"]
+
+    @classmethod
+    def can_handle(
+        cls,
+        model: model_types.SupportedModelType,
+    ) -> TypeGuard["keras.Model"]:
+        if not type_utils.LazyType("keras.Model").isinstance(model):
+            return False
+        import keras
+
+        return version.parse(keras.__version__) >= version.parse("3.0.0")
+
+    @classmethod
+    def cast_model(
+        cls,
+        model: model_types.SupportedModelType,
+    ) -> "keras.Model":
+        import keras
+
+        assert isinstance(model, keras.Model)
+
+        return cast(keras.Model, model)
+
+    @classmethod
+    def save_model(
+        cls,
+        name: str,
+        model: "keras.Model",
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
+        **kwargs: Unpack[model_types.TensorflowSaveOptions],
+    ) -> None:
+        enable_explainability = kwargs.get("enable_explainability", False)
+        if enable_explainability:
+            raise NotImplementedError("Explainability is not supported for Tensorflow model.")
+
+        import keras
+
+        assert isinstance(model, keras.Model)
+
+        if not is_sub_model:
+            target_methods = handlers_utils.get_target_methods(
+                model=model,
+                target_methods=kwargs.pop("target_methods", None),
+                default_target_methods=cls.DEFAULT_TARGET_METHODS,
+            )
+
+            def get_prediction(
+                target_method_name: str, sample_input_data: "model_types.SupportedLocalDataType"
+            ) -> model_types.SupportedLocalDataType:
+                target_method = getattr(model, target_method_name, None)
+                assert callable(target_method)
+                predictions_df = target_method(sample_input_data)
+
+                if (
+                    type_utils.LazyType("tensorflow.Tensor").isinstance(predictions_df)
+                    or type_utils.LazyType("tensorflow.Variable").isinstance(predictions_df)
+                    or type_utils.LazyType("torch.Tensor").isinstance(predictions_df)
+                ):
+                    predictions_df = [predictions_df]
+
+                return predictions_df
+
+            model_meta = handlers_utils.validate_signature(
+                model=model,
+                model_meta=model_meta,
+                target_methods=target_methods,
+                sample_input_data=sample_input_data,
+                get_prediction_fn=get_prediction,
+            )
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        os.makedirs(model_blob_path, exist_ok=True)
+        save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+        model.save(save_path)
+
+        custom_object_save_path = os.path.join(model_blob_path, cls.CUSTOM_OBJECT_SAVE_PATH)
+        custom_objects = keras.saving.get_custom_objects()
+        with open(custom_object_save_path, "wb") as f:
+            cloudpickle.dump(custom_objects, f)
+
+        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,
+        )
+        model_meta.models[name] = base_meta
+        model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
+
+        dependencies = [
+            model_env.ModelDependency(requirement="keras>=3", pip_name="keras"),
+        ]
+        keras_backend = keras.backend.backend()
+        if keras_backend == "tensorflow":
+            dependencies.append(model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"))
+        elif keras_backend == "torch":
+            dependencies.append(model_env.ModelDependency(requirement="pytorch", pip_name="torch"))
+        elif keras_backend == "jax":
+            dependencies.append(model_env.ModelDependency(requirement="jax", pip_name="jax"))
+        else:
+            raise ValueError(f"Unsupported backend {keras_backend}")
+
+        model_meta.env.include_if_absent(
+            dependencies,
+            check_local_version=True,
+        )
+        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+
+    @classmethod
+    def load_model(
+        cls,
+        name: str,
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        **kwargs: Unpack[model_types.TensorflowLoadOptions],
+    ) -> "keras.Model":
+        import keras
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        model_blobs_metadata = model_meta.models
+        model_blob_metadata = model_blobs_metadata[name]
+        model_blob_filename = model_blob_metadata.path
+
+        custom_object_save_path = os.path.join(model_blob_path, cls.CUSTOM_OBJECT_SAVE_PATH)
+        with open(custom_object_save_path, "rb") as f:
+            custom_objects = cloudpickle.load(f)
+        load_path = os.path.join(model_blob_path, model_blob_filename)
+        m = keras.models.load_model(load_path, custom_objects=custom_objects, safe_mode=False)
+
+        return cast(keras.Model, m)
+
+    @classmethod
+    def convert_as_custom_model(
+        cls,
+        raw_model: "keras.Model",
+        model_meta: model_meta_api.ModelMetadata,
+        background_data: Optional[pd.DataFrame] = None,
+        **kwargs: Unpack[model_types.TensorflowLoadOptions],
+    ) -> custom_model.CustomModel:
+
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: "keras.Model",
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: "keras.Model",
+                signature: model_signature.ModelSignature,
+                target_method: str,
+            ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
+                dtype_map = {
+                    spec.name: spec.as_dtype(force_numpy_dtype=True)
+                    for spec in signature.inputs
+                    if isinstance(spec, model_signature.FeatureSpec)
+                }
+
+                @custom_model.inference_api
+                def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    res = getattr(raw_model, target_method)(X.astype(dtype_map), verbose=0)
+
+                    if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
+                        # In case of multi-output estimators, predict_proba(), decision_function(), etc., functions
+                        # return a list of ndarrays. We need to deal them separately
+                        df = numpy_handler.SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                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)
+
+            _KerasModel = type(
+                "_KerasModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _KerasModel
+
+        _KerasModel = _create_custom_model(raw_model, model_meta)
+        keras_model = _KerasModel(custom_model.ModelContext())
+
+        return keras_model

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

@@ -49,6 +49,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
             type_utils.LazyType("torch.nn.Module").isinstance(model)
             and not type_utils.LazyType("torch.jit.ScriptModule").isinstance(model)
             and not type_utils.LazyType("sentence_transformers.SentenceTransformer").isinstance(model)
+            and not type_utils.LazyType("keras.Model").isinstance(model)
         )
 
     @classmethod

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

@@ -292,12 +292,37 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
                     import shap
 
-                    # TODO: if not resolved by explainer, we need to pass the callable function
                     try:
                         explainer = shap.Explainer(raw_model, background_data)
                         df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
-                    except TypeError as e:
-                        raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
+                    except TypeError:
+                        try:
+                            dtype_map = {
+                                spec.name: spec.as_dtype(force_numpy_dtype=True)  # type: ignore[attr-defined]
+                                for spec in signature.inputs
+                            }
+
+                            if isinstance(X, pd.DataFrame):
+                                X = X.astype(dtype_map, copy=False)
+                            if hasattr(raw_model, "predict_proba"):
+                                if isinstance(X, np.ndarray):
+                                    explanations = shap.Explainer(
+                                        raw_model.predict_proba, background_data.values  # type: ignore[union-attr]
+                                    )(X).values
+                                else:
+                                    explanations = shap.Explainer(raw_model.predict_proba, background_data)(X).values
+                            elif hasattr(raw_model, "predict"):
+                                if isinstance(X, np.ndarray):
+                                    explanations = shap.Explainer(
+                                        raw_model.predict, background_data.values  # type: ignore[union-attr]
+                                    )(X).values
+                                else:
+                                    explanations = shap.Explainer(raw_model.predict, background_data)(X).values
+                            else:
+                                raise ValueError("Missing any supported target method to explain.")
+                            df = handlers_utils.convert_explanations_to_2D_df(raw_model, explanations)
+                        except TypeError as e:
+                            raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
                 if target_method == "explain":

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

@@ -74,11 +74,6 @@ class SnowMLModelHandler(_base.BaseModelHandler["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
 
         tree_methods = ["to_xgboost", "to_lightgbm", "to_sklearn"]
         non_tree_methods = ["to_sklearn"]
@@ -129,27 +124,54 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         # Pipeline is inherited from BaseEstimator, so no need to add one more check
 
         if not is_sub_model:
-            if model_meta.signatures:
+            if model_meta.signatures or sample_input_data is not None:
                 warnings.warn(
                     "Providing model signature for Snowpark ML "
                     + "Modeling model is not required. Model signature will automatically be inferred during fitting. ",
                     UserWarning,
                     stacklevel=2,
                 )
-            assert hasattr(model, "model_signatures"), "Model does not have model signatures as expected."
-            model_signature_dict = getattr(model, "model_signatures", {})
-            target_methods = kwargs.pop("target_methods", None)
-            if not target_methods:
-                model_meta.signatures = model_signature_dict
+                target_methods = handlers_utils.get_target_methods(
+                    model=model,
+                    target_methods=kwargs.pop("target_methods", None),
+                    default_target_methods=cls.DEFAULT_TARGET_METHODS,
+                )
+
+                def get_prediction(
+                    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)
+
+                    target_method = getattr(model, target_method_name, None)
+                    assert callable(target_method)
+                    predictions_df = target_method(sample_input_data)
+                    return predictions_df
+
+                model_meta = handlers_utils.validate_signature(
+                    model=model,
+                    model_meta=model_meta,
+                    target_methods=target_methods,
+                    sample_input_data=sample_input_data,
+                    get_prediction_fn=get_prediction,
+                    is_for_modeling_model=True,
+                )
             else:
-                temp_model_signature_dict = {}
-                for method_name in target_methods:
-                    method_model_signature = model_signature_dict.get(method_name, None)
-                    if method_model_signature is not None:
-                        temp_model_signature_dict[method_name] = method_model_signature
-                    else:
-                        raise ValueError(f"Target method {method_name} does not exist in the model.")
-                model_meta.signatures = temp_model_signature_dict
+                assert hasattr(model, "model_signatures"), "Model does not have model signatures as expected."
+                model_signature_dict = getattr(model, "model_signatures", {})
+                optional_target_methods = kwargs.pop("target_methods", None)
+                if not optional_target_methods:
+                    model_meta.signatures = model_signature_dict
+                else:
+                    temp_model_signature_dict = {}
+                    for method_name in optional_target_methods:
+                        method_model_signature = model_signature_dict.get(method_name, None)
+                        if method_model_signature is not None:
+                            temp_model_signature_dict[method_name] = method_model_signature
+                        else:
+                            raise ValueError(f"Target method {method_name} does not exist in the model.")
+                    model_meta.signatures = temp_model_signature_dict
 
         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)
@@ -279,9 +301,40 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                     for method_name in non_tree_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)
+                            try:
+                                explainer = shap.Explainer(base_model, masker=background_data)
+                                df = handlers_utils.convert_explanations_to_2D_df(base_model, explainer(X).values)
+                            except TypeError:
+                                try:
+                                    dtype_map = {
+                                        spec.name: spec.as_dtype(force_numpy_dtype=True)  # type: ignore[attr-defined]
+                                        for spec in signature.inputs
+                                    }
+
+                                    if isinstance(X, pd.DataFrame):
+                                        X = X.astype(dtype_map, copy=False)
+                                    if hasattr(base_model, "predict_proba"):
+                                        if isinstance(X, np.ndarray):
+                                            explainer = shap.Explainer(
+                                                base_model.predict_proba,
+                                                background_data.values,  # type: ignore[union-attr]
+                                            )
+                                        else:
+                                            explainer = shap.Explainer(base_model.predict_proba, background_data)
+                                    elif hasattr(base_model, "predict"):
+                                        if isinstance(X, np.ndarray):
+                                            explainer = shap.Explainer(
+                                                base_model.predict, background_data.values  # type: ignore[union-attr]
+                                            )
+                                        else:
+                                            explainer = shap.Explainer(base_model.predict, background_data)
+                                    else:
+                                        raise ValueError("Missing any supported target method to explain.")
+                                    df = handlers_utils.convert_explanations_to_2D_df(base_model, explainer(X).values)
+                                except TypeError as e:
+                                    raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
                             return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
                         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.")

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

@@ -10,7 +10,10 @@ 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_migrator import base_migrator
+from snowflake.ml.model._packager.model_handlers_migrator import (
+    base_migrator,
+    tensorflow_migrator_2023_12_01,
+)
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
@@ -28,15 +31,17 @@ if TYPE_CHECKING:
 
 @final
 class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
-    """Handler for TensorFlow based model.
+    """Handler for TensorFlow based model or keras v2 model.
 
     Currently tensorflow.Module based classes are supported.
     """
 
     HANDLER_TYPE = "tensorflow"
-    HANDLER_VERSION = "2023-12-01"
-    _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    HANDLER_VERSION = "2025-01-01"
+    _MIN_SNOWPARK_ML_VERSION = "1.7.5"
+    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+        "2023-12-01": tensorflow_migrator_2023_12_01.TensorflowHandlerMigrator20231201
+    }
 
     MODEL_BLOB_FILE_OR_DIR = "model"
     DEFAULT_TARGET_METHODS = ["__call__"]
@@ -46,7 +51,13 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         cls,
         model: model_types.SupportedModelType,
     ) -> TypeGuard["tensorflow.nn.Module"]:
-        return type_utils.LazyType("tensorflow.Module").isinstance(model)
+        if not type_utils.LazyType("tensorflow.Module").isinstance(model):
+            return False
+        if type_utils.LazyType("keras.Model").isinstance(model):
+            import keras
+
+            return version.parse(keras.__version__) < version.parse("3.0.0")
+        return True
 
     @classmethod
     def cast_model(
@@ -74,44 +85,22 @@ 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(
-            "keras.Model"
-        ).isinstance(model)
+        is_keras_model = type_utils.LazyType("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
+        # Keras v2 models are saved using keras api
+        # Tensorflow models are saved using tensorflow api
 
         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"
+            save_format = "keras_tf"
         else:
             save_format = "tf"
 
-        if is_keras_model:
+        if is_keras_model or is_tf_keras_model:
             default_target_methods = ["predict"]
         else:
             default_target_methods = cls.DEFAULT_TARGET_METHODS
@@ -156,15 +145,8 @@ 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)
         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":
+        if 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,
@@ -186,7 +168,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"),
         ]
         if is_keras_model:
-            dependencies.append(model_env.ModelDependency(requirement="keras", pip_name="keras"))
+            dependencies.append(model_env.ModelDependency(requirement="keras<=3", pip_name="keras"))
         elif is_tf_keras_model:
             dependencies.append(model_env.ModelDependency(requirement="tf-keras", pip_name="tf-keras"))
 
@@ -204,6 +186,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blobs_dir_path: str,
         **kwargs: Unpack[model_types.TensorflowLoadOptions],
     ) -> "tensorflow.Module":
+        os.environ["TF_USE_LEGACY_KERAS"] = "1"
         import tensorflow
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
@@ -212,14 +195,9 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blob_filename = model_blob_metadata.path
         model_blob_options = cast(model_meta_schema.TensorflowModelBlobOptions, model_blob_metadata.options)
         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":
+        save_format = model_blob_options.get("save_format", "keras_tf")
+        if 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(load_path)
 

snowflake/ml/model/_packager/model_handlers_migrator/tensorflow_migrator_2023_12_01.py

@@ -0,0 +1,48 @@
+from typing import cast
+
+from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
+from snowflake.ml.model._packager.model_meta import (
+    model_meta as model_meta_api,
+    model_meta_schema,
+)
+
+
+class TensorflowHandlerMigrator20231201(base_migrator.BaseModelHandlerMigrator):
+    source_version = "2023-12-01"
+    target_version = "2025-01-01"
+
+    @staticmethod
+    def upgrade(name: str, model_meta: model_meta_api.ModelMetadata, model_blobs_dir_path: str) -> None:
+
+        model_blob_metadata = model_meta.models[name]
+        model_blob_options = cast(model_meta_schema.TensorflowModelBlobOptions, model_blob_metadata.options)
+        # To migrate code <= 1.7.0, default to keras model
+        is_old_model = "save_format" not in model_blob_options and "is_keras_model" not in model_blob_options
+        # To migrate code form 1.7.1, default to False.
+        is_keras_model = model_blob_options.get("is_keras_model", False)
+        # To migrate code from 1.7.2, default to tf, has options keras, keras_tf, cloudpickle, tf
+        #
+        # 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"
+        #
+        save_format = model_blob_options.get("save_format", "tf")
+
+        if save_format == "keras" or is_keras_model or is_old_model:
+            save_format = "keras_tf"
+        elif save_format == "cloudpickle":
+            # Given the old logic, this could only happen if the original model is a keras model, and keras is 3.x
+            # However, in this case, keras.Model does not extends from tensorflow.Module
+            # So actually TensorflowHandler will not be triggered, we could safely error this out.
+            raise NotImplementedError(
+                "Unable to upgrade keras 3.x model saved by old handler. This is not supposed to happen"
+            )
+
+        model_blob_options["save_format"] = save_format
+        model_meta.models[name].options = model_blob_options

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

@@ -352,7 +352,7 @@ class ModelMetadata:
             version.parse(loaded_meta_min_snowpark_ml_version) > version.parse(snowml_env.VERSION)
         ):
             raise RuntimeError(
-                f"The minimal version required to load the model is {loaded_meta_min_snowpark_ml_version},"
+                f"The minimal version required to load the model is {loaded_meta_min_snowpark_ml_version}, "
                 f"while current version of Snowpark ML library is {snowml_env.VERSION}."
             )
         return model_meta_schema.ModelMetadataDict(

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

@@ -44,6 +44,9 @@ class CatBoostModelBlobOptions(BaseModelBlobOptions):
 class HuggingFacePipelineModelBlobOptions(BaseModelBlobOptions):
     task: Required[str]
     batch_size: Required[int]
+    has_tokenizer: NotRequired[bool]
+    has_feature_extractor: NotRequired[bool]
+    has_image_preprocessor: NotRequired[bool]
 
 
 class LightGBMModelBlobOptions(BaseModelBlobOptions):

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,<5', '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,!=1.26.0', '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,<5', '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,!=1.26.0', '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']
+ALL_REQUIREMENTS=['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<5', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2024.6.1,<2026', 'importlib_resources>=6.1.1, <7', 'keras>=2.0.0,<4', '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,<3', '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.7.0,<3', 'sentencepiece>=0.1.95,<0.2.0', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2,!=1.26.0', 'sqlparse>=0.4,<1', 'tensorflow>=2.17.0,<3', 'tokenizers>=0.15.1,<1', 'torchdata>=0.4,<1', 'transformers>=4.37.2,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']

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

@@ -45,6 +45,7 @@ class ModelRuntime:
         self.name = name
         self.runtime_env = copy.deepcopy(env)
         self.imports = imports or []
+        self.is_gpu = is_gpu
 
         if loading_from_file:
             return
@@ -88,7 +89,9 @@ class ModelRuntime:
         self.runtime_env.conda_env_rel_path = self.runtime_rel_path / self.runtime_env.conda_env_rel_path
         self.runtime_env.pip_requirements_rel_path = self.runtime_rel_path / self.runtime_env.pip_requirements_rel_path
 
-        env_dict = self.runtime_env.save_as_dict(packager_path, default_channel_override=default_channel_override)
+        env_dict = self.runtime_env.save_as_dict(
+            packager_path, default_channel_override=default_channel_override, is_gpu=self.is_gpu
+        )
 
         return model_meta_schema.ModelRuntimeDict(
             imports=list(map(str, self.imports)),