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

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -27,7 +27,7 @@ def get_model_method_options_from_options(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
     default_function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
-    if options.get("enable_explainability", False) and target_method.startswith("explain"):
+    if target_method == "explain":
         default_function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
     method_option = options.get("method_options", {}).get(target_method, {})
     global_function_type = options.get("function_type", default_function_type)

snowflake/ml/model/_packager/model_env/model_env.py

@@ -174,6 +174,18 @@ class ModelEnv:
             except env_utils.DuplicateDependencyError:
                 pass
 
+    def remove_if_present_conda(self, conda_pkgs: List[str]) -> None:
+        """Remove conda requirements from model env if present.
+
+        Args:
+            conda_pkgs: A list of package name to be removed from conda requirements.
+        """
+        for pkg_name in conda_pkgs:
+            spec_conda = env_utils._find_conda_dep_spec(self._conda_dependencies, pkg_name)
+            if spec_conda:
+                channel, spec = spec_conda
+                self._conda_dependencies[channel].remove(spec)
+
     def generate_env_for_cuda(self) -> None:
         if self.cuda_version is None:
             return

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

@@ -179,7 +179,7 @@ def convert_explanations_to_2D_df(
         return pd.DataFrame(explanations)
 
     if hasattr(model, "classes_"):
-        classes_list = [str(cl) for cl in model.classes_]  # type:ignore[union-attr]
+        classes_list = [str(cl) for cl in model.classes_]
         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]}")
@@ -191,7 +191,11 @@ def convert_explanations_to_2D_df(
         # convert to object or numpy creates strings of fixed length
         return np.asarray(json.dumps(dict(zip(classes_list, row)), cls=NumpyEncoder), dtype=object)
 
-    exp_2d = np.apply_along_axis(row_to_dict, -1, explanations)
+    # convert to dict only for multiclass
+    if len(classes_list) > 2:
+        exp_2d = np.apply_along_axis(row_to_dict, -1, explanations)
+    else:  # assumes index 1 is positive class always
+        exp_2d = np.apply_along_axis(lambda arr: arr[1], -1, explanations)
 
     return pd.DataFrame(exp_2d)
 

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

@@ -9,17 +9,14 @@ 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,
-    model_objective_utils,
-)
+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,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -97,8 +94,8 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
-            model_meta.task = model_task_and_output.task
+            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)
             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(

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

@@ -2,7 +2,7 @@ import inspect
 import os
 import pathlib
 import sys
-from typing import Dict, Optional, Type, final
+from typing import Dict, Optional, Type, cast, final
 
 import anyio
 import cloudpickle
@@ -99,6 +99,8 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
             for sub_name, model_ref in model.context.model_refs.items():
                 handler = model_handler.find_handler(model_ref.model)
                 assert handler is not None
+                if handler is None:
+                    raise TypeError("Your input type to custom model is not currently supported")
                 sub_model = handler.cast_model(model_ref.model)
                 handler.save_model(
                     name=sub_name,
@@ -106,6 +108,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                     model_meta=model_meta,
                     model_blobs_dir_path=model_blobs_dir_path,
                     is_sub_model=True,
+                    **cast(model_types.BaseModelSaveOption, kwargs),
                 )
 
         # Make sure that the module where the model is defined get pickled by value as well.
@@ -173,6 +176,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 name=sub_model_name,
                 model_meta=model_meta,
                 model_blobs_dir_path=model_blobs_dir_path,
+                **cast(model_types.BaseModelLoadOption, kwargs),
             )
             models[sub_model_name] = sub_model
         reconstructed_context = custom_model.ModelContext(artifacts=artifacts, models=models)

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

@@ -256,12 +256,20 @@ class HuggingFacePipelineHandler(
     @staticmethod
     def _get_device_config(**kwargs: Unpack[model_types.HuggingFaceLoadOptions]) -> Dict[str, str]:
         device_config: Dict[str, Any] = {}
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+        gpu_nums = 0
+        if cuda_visible_devices is not None:
+            gpu_nums = len(cuda_visible_devices.split(","))
         if (
             kwargs.get("use_gpu", False)
             and kwargs.get("device_map", None) is None
             and kwargs.get("device", None) is None
         ):
-            device_config["device_map"] = "auto"
+            if gpu_nums == 0 or gpu_nums > 1:
+                # Use accelerator if there are multiple GPUs or no GPU
+                device_config["device_map"] = "auto"
+            else:
+                device_config["device"] = "cuda"
         elif kwargs.get("device_map", None) is not None:
             device_config["device_map"] = kwargs["device_map"]
         elif kwargs.get("device", None) is not None:
@@ -310,6 +318,7 @@ class HuggingFacePipelineHandler(
             m = transformers.pipeline(
                 model_blob_options["task"],
                 model=model_blob_file_or_dir_path,
+                trust_remote_code=True,
                 **device_config,
             )
 

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

@@ -20,17 +20,14 @@ 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,
-    model_objective_utils,
-)
+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,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -113,7 +110,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
+            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)
             if enable_explainability:
                 explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
@@ -199,13 +196,14 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         with open(model_blob_file_path, "rb") as f:
             model = cloudpickle.load(f)
         assert isinstance(model, getattr(lightgbm, lightgbm_estimator_type))
+        assert isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel)
 
         return model
 
     @classmethod
     def convert_as_custom_model(
         cls,
-        raw_model: Union["lightgbm.Booster", "lightgbm.XGBModel"],
+        raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
         model_meta: model_meta_api.ModelMetadata,
         background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.LGBMModelLoadOptions],

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

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import os
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
@@ -155,8 +156,14 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         model_blob_filename = model_blob_metadata.path
         model_blob_file_or_dir_path = os.path.join(model_blob_path, model_blob_filename)
 
+        additional_kwargs = {}
+        if "trust_remote_code" in inspect.signature(sentence_transformers.SentenceTransformer).parameters:
+            additional_kwargs["trust_remote_code"] = True
+
         model = sentence_transformers.SentenceTransformer(
-            model_blob_file_or_dir_path, device=cls._get_device_config(**kwargs)
+            model_blob_file_or_dir_path,
+            device=cls._get_device_config(**kwargs),
+            **additional_kwargs,
         )
         return model
 

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

@@ -10,24 +10,35 @@ 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,
-    model_objective_utils,
-)
+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,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
+from snowflake.ml.modeling._internal.constants import IN_ML_RUNTIME_ENV_VAR
 
 if TYPE_CHECKING:
     import sklearn.base
     import sklearn.pipeline
 
 
+def _unpack_container_runtime_pipeline(model: "sklearn.pipeline.Pipeline") -> "sklearn.pipeline.Pipeline":
+    new_steps = []
+    for step_name, step in model.steps:
+        new_reg = step
+        if hasattr(step, "_sklearn_estimator") and step._sklearn_estimator is not None:
+            # Unpack estimator to open source.
+            new_reg = step._sklearn_estimator
+        new_steps.append((step_name, new_reg))
+
+    model.steps = new_steps
+    return model
+
+
 @final
 class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]]):
     """Handler for scikit-learn based model.
@@ -104,6 +115,10 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             if sample_input_data is None:
                 raise ValueError("Sample input data is required to enable explainability.")
 
+        # If this is a pipeline and we are in the container runtime, check for distributed estimator.
+        if os.getenv(IN_ML_RUNTIME_ENV_VAR) and isinstance(model, sklearn.pipeline.Pipeline):
+            model = _unpack_container_runtime_pipeline(model)
+
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
                 model=model,
@@ -137,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_objective_utils.get_model_task_and_output_type(model)
-            model_meta.task = model_task_and_output_type.task
+            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)
 
             # if users did not ask then we enable if we have background data
             if enable_explainability is None:
@@ -180,6 +195,35 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
+        # if model instance is a pipeline, check the pipeline steps
+        if isinstance(model, sklearn.pipeline.Pipeline):
+            for _, pipeline_step in model.steps:
+                if type_utils.LazyType("lightgbm.LGBMModel").isinstance(pipeline_step) or type_utils.LazyType(
+                    "lightgbm.Booster"
+                ).isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="lightgbm", pip_name="lightgbm"),
+                        ],
+                        check_local_version=True,
+                    )
+                elif type_utils.LazyType("xgboost.XGBModel").isinstance(pipeline_step) or type_utils.LazyType(
+                    "xgboost.Booster"
+                ).isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
+                        ],
+                        check_local_version=True,
+                    )
+                elif type_utils.LazyType("catboost.CatBoost").isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="catboost", pip_name="catboost"),
+                        ],
+                        check_local_version=True,
+                    )
+
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP

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

@@ -5,24 +5,20 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, fin
 import cloudpickle
 import numpy as np
 import pandas as pd
-from packaging import version
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import type_utils
 from snowflake.ml._internal.exceptions import exceptions
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_handlers import (
-    _base,
-    _utils as handlers_utils,
-    model_objective_utils,
-)
+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,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -72,41 +68,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         return cast("BaseEstimator", model)
 
     @classmethod
-    def _get_local_version_package(cls, pkg_name: str) -> Optional[version.Version]:
-        from importlib import metadata as importlib_metadata
-
-        from packaging import version
-
-        local_version = None
-
-        try:
-            local_dist = importlib_metadata.distribution(pkg_name)
-            local_version = version.parse(local_dist.version)
-        except importlib_metadata.PackageNotFoundError:
-            pass
-
-        return local_version
-
-    @classmethod
-    def _can_support_xgb(cls, enable_explainability: Optional[bool]) -> bool:
-
-        local_xgb_version = cls._get_local_version_package("xgboost")
-
-        if local_xgb_version and local_xgb_version >= version.parse("2.1.0"):
-            if enable_explainability:
-                warnings.warn(
-                    f"This version of xgboost {local_xgb_version} does not work with shap 0.42.1."
-                    + "If you want model explanations, lower the xgboost version to <2.1.0.",
-                    category=UserWarning,
-                    stacklevel=1,
-                )
-            return False
-        return True
-
-    @classmethod
-    def _get_supported_object_for_explainability(
-        cls, estimator: "BaseEstimator", enable_explainability: Optional[bool]
-    ) -> Any:
+    def _get_supported_object_for_explainability(cls, estimator: "BaseEstimator") -> Any:
         from snowflake.ml.modeling import pipeline as snowml_pipeline
 
         # handle pipeline objects separately
@@ -118,8 +80,6 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
             if hasattr(estimator, method_name):
                 try:
                     result = getattr(estimator, method_name)()
-                    if method_name == "to_xgboost" and not cls._can_support_xgb(enable_explainability):
-                        return None
                     return result
                 except exceptions.SnowflakeMLException:
                     pass  # Do nothing and continue to the next method
@@ -168,7 +128,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 model_meta.signatures = temp_model_signature_dict
 
         if enable_explainability or enable_explainability is None:
-            python_base_obj = cls._get_supported_object_for_explainability(model, enable_explainability)
+            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(
@@ -177,8 +137,8 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 # set None to False so we don't include shap in the environment
                 enable_explainability = False
             else:
-                model_task_and_output_type = model_objective_utils.get_model_task_and_output_type(python_base_obj)
-                model_meta.task = model_task_and_output_type.task
+                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,
@@ -213,28 +173,10 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         model_dependencies = model._get_dependencies()
         for dep in model_dependencies:
             pkg_name = dep.split("==")[0]
-            if pkg_name != "xgboost":
-                _include_if_absent_pkgs.append(model_env.ModelDependency(requirement=pkg_name, pip_name=pkg_name))
-                continue
-
-            local_xgb_version = cls._get_local_version_package("xgboost")
-            if local_xgb_version and local_xgb_version >= version.parse("2.0.0") and enable_explainability:
-                model_meta.env.include_if_absent(
-                    [
-                        model_env.ModelDependency(requirement="xgboost==2.0.*", pip_name="xgboost"),
-                    ],
-                    check_local_version=False,
-                )
-            else:
-                model_meta.env.include_if_absent(
-                    [
-                        model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
-                    ],
-                    check_local_version=True,
-                )
+            _include_if_absent_pkgs.append(model_env.ModelDependency(requirement=pkg_name, pip_name=pkg_name))
 
         if enable_explainability:
-            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
+            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap>=0.46.0", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
         model_meta.env.include_if_absent(_include_if_absent_pkgs, check_local_version=True)
 

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

@@ -13,6 +13,7 @@ from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
+    model_meta_schema,
 )
 from snowflake.ml.model._signatures import (
     numpy_handler,
@@ -76,7 +77,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
 
         assert isinstance(model, tensorflow.Module)
 
-        if isinstance(model, tensorflow.keras.Model):
+        is_keras_model = type_utils.LazyType("tensorflow.keras.Model").isinstance(model) or type_utils.LazyType(
+            "tf_keras.Model"
+        ).isinstance(model)
+
+        if is_keras_model:
             default_target_methods = ["predict"]
         else:
             default_target_methods = cls.DEFAULT_TARGET_METHODS
@@ -117,8 +122,14 @@ 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 isinstance(model, tensorflow.keras.Model):
+        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,
+            )
         else:
             tensorflow.saved_model.save(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
 
@@ -127,12 +138,16 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             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),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
         model_meta.env.include_if_absent(
-            [model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow")], check_local_version=True
+            [
+                model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"),
+            ],
+            check_local_version=True,
         )
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
@@ -150,9 +165,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blobs_metadata = model_meta.models
         model_blob_metadata = model_blobs_metadata[name]
         model_blob_filename = model_blob_metadata.path
-        m = tensorflow.keras.models.load_model(os.path.join(model_blob_path, model_blob_filename), compile=False)
-        if isinstance(m, tensorflow.keras.Model):
-            return m
+        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)
+        else:
+            m = tensorflow.saved_model.load(os.path.join(model_blob_path, model_blob_filename))
         return cast(tensorflow.Module, m)
 
     @classmethod

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

@@ -23,7 +23,7 @@ if TYPE_CHECKING:
 
 
 @final
-class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # type:ignore[name-defined]
+class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     """Handler for PyTorch JIT based model.
 
     Currently torch.jit.ScriptModule based classes are supported.
@@ -41,25 +41,25 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
     def can_handle(
         cls,
         model: model_types.SupportedModelType,
-    ) -> TypeGuard["torch.jit.ScriptModule"]:  # type:ignore[name-defined]
+    ) -> TypeGuard["torch.jit.ScriptModule"]:
         return type_utils.LazyType("torch.jit.ScriptModule").isinstance(model)
 
     @classmethod
     def cast_model(
         cls,
         model: model_types.SupportedModelType,
-    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+    ) -> "torch.jit.ScriptModule":
         import torch
 
-        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(model, torch.jit.ScriptModule)
 
-        return cast(torch.jit.ScriptModule, model)  # type:ignore[name-defined]
+        return cast(torch.jit.ScriptModule, model)
 
     @classmethod
     def save_model(
         cls,
         name: str,
-        model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        model: "torch.jit.ScriptModule",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
@@ -72,7 +72,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
 
         import torch
 
-        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(model, torch.jit.ScriptModule)
 
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -111,7 +111,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
         with open(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
-            torch.jit.save(model, f)  # type:ignore[no-untyped-call, attr-defined]
+            torch.jit.save(model, f)  # type:ignore[no-untyped-call]
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
@@ -133,7 +133,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         **kwargs: Unpack[model_types.TorchScriptLoadOptions],
-    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+    ) -> "torch.jit.ScriptModule":
         import torch
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
@@ -141,10 +141,10 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_blob_metadata = model_blobs_metadata[name]
         model_blob_filename = model_blob_metadata.path
         with open(os.path.join(model_blob_path, model_blob_filename), "rb") as f:
-            m = torch.jit.load(  # type:ignore[no-untyped-call, attr-defined]
+            m = torch.jit.load(  # type:ignore[no-untyped-call]
                 f, map_location="cuda" if kwargs.get("use_gpu", False) else "cpu"
             )
-        assert isinstance(m, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(m, torch.jit.ScriptModule)
 
         if kwargs.get("use_gpu", False):
             m = m.cuda()
@@ -154,7 +154,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
     @classmethod
     def convert_as_custom_model(
         cls,
-        raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        raw_model: "torch.jit.ScriptModule",
         model_meta: model_meta_api.ModelMetadata,
         background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.TorchScriptLoadOptions],
@@ -162,11 +162,11 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         from snowflake.ml.model import custom_model
 
         def _create_custom_model(
-            raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+            raw_model: "torch.jit.ScriptModule",
             model_meta: model_meta_api.ModelMetadata,
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
-                raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+                raw_model: "torch.jit.ScriptModule",
                 signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]: