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

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

@@ -8,10 +8,14 @@ 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,
+    torchscript_migrator_2023_12_01,
+)
 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 (
     pytorch_handler,
@@ -30,9 +34,11 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     """
 
     HANDLER_TYPE = "torchscript"
-    HANDLER_VERSION = "2023-12-01"
-    _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    HANDLER_VERSION = "2025-03-01"
+    _MIN_SNOWPARK_ML_VERSION = "1.8.0"
+    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+        "2023-12-01": torchscript_migrator_2023_12_01.TorchScriptHandlerMigrator20231201
+    }
 
     MODEL_BLOB_FILE_OR_DIR = "model.pt"
     DEFAULT_TARGET_METHODS = ["forward"]
@@ -81,22 +87,32 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
                 default_target_methods=cls.DEFAULT_TARGET_METHODS,
             )
 
+            multiple_inputs = kwargs.get("multiple_inputs", False)
+
             def get_prediction(
                 target_method_name: str, sample_input_data: "model_types.SupportedLocalDataType"
             ) -> model_types.SupportedLocalDataType:
-                if not pytorch_handler.SeqOfPyTorchTensorHandler.can_handle(sample_input_data):
-                    sample_input_data = pytorch_handler.SeqOfPyTorchTensorHandler.convert_from_df(
-                        model_signature._convert_local_data_to_df(sample_input_data)
-                    )
+                if multiple_inputs:
+                    if not pytorch_handler.SeqOfPyTorchTensorHandler.can_handle(sample_input_data):
+                        sample_input_data = pytorch_handler.SeqOfPyTorchTensorHandler.convert_from_df(
+                            model_signature._convert_local_data_to_df(sample_input_data)
+                        )
+                else:
+                    if not pytorch_handler.PyTorchTensorHandler.can_handle(sample_input_data):
+                        sample_input_data = pytorch_handler.PyTorchTensorHandler.convert_from_df(
+                            model_signature._convert_local_data_to_df(sample_input_data)
+                        )
 
                 model.eval()
                 target_method = getattr(model, target_method_name, None)
                 assert callable(target_method)
                 with torch.no_grad():
-                    predictions_df = target_method(*sample_input_data)
-
-                if isinstance(predictions_df, torch.Tensor):
-                    predictions_df = [predictions_df]
+                    if multiple_inputs:
+                        predictions_df = target_method(*sample_input_data)
+                        if not isinstance(predictions_df, tuple):
+                            predictions_df = [predictions_df]
+                    else:
+                        predictions_df = target_method(sample_input_data)
 
                 return predictions_df
 
@@ -117,6 +133,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
+            options=model_meta_schema.TorchScriptModelBlobOptions(multiple_inputs=multiple_inputs),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -170,6 +187,10 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
                 signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
+                multiple_inputs = cast(
+                    model_meta_schema.TorchScriptModelBlobOptions, model_meta.models[model_meta.name].options
+                )["multiple_inputs"]
+
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
                     if X.isnull().any(axis=None):
@@ -179,19 +200,27 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
 
                     raw_model.eval()
 
-                    t = pytorch_handler.SeqOfPyTorchTensorHandler.convert_from_df(X, signature.inputs)
+                    if multiple_inputs:
+                        st = pytorch_handler.SeqOfPyTorchTensorHandler.convert_from_df(X, signature.inputs)
 
-                    if kwargs.get("use_gpu", False):
-                        t = [element.cuda() for element in t]
+                        if kwargs.get("use_gpu", False):
+                            st = [element.cuda() for element in st]
 
-                    with torch.no_grad():
-                        res = getattr(raw_model, target_method)(*t)
+                        with torch.no_grad():
+                            res = getattr(raw_model, target_method)(*st)
 
-                    if isinstance(res, torch.Tensor):
-                        res = [res]
+                        if not isinstance(res, tuple):
+                            res = [res]
+                    else:
+                        t = pytorch_handler.PyTorchTensorHandler.convert_from_df(X, signature.inputs)
+                        if kwargs.get("use_gpu", False):
+                            t = t.cuda()
 
+                        with torch.no_grad():
+                            res = getattr(raw_model, target_method)(t)
                     return model_signature_utils.rename_pandas_df(
-                        data=pytorch_handler.SeqOfPyTorchTensorHandler.convert_to_df(res), features=signature.outputs
+                        model_signature._convert_local_data_to_df(res, ensure_serializable=True),
+                        features=signature.outputs,
                     )
 
                 return fn

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

@@ -99,10 +99,10 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
             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)):
+                if not isinstance(sample_input_data, (pd.DataFrame, np.ndarray, xgboost.DMatrix)):
                     sample_input_data = model_signature._convert_local_data_to_df(sample_input_data)
 
-                if isinstance(model, xgboost.Booster):
+                if isinstance(model, xgboost.Booster) and not isinstance(sample_input_data, xgboost.DMatrix):
                     sample_input_data = xgboost.DMatrix(sample_input_data)
 
                 target_method = getattr(model, target_method_name, None)

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

@@ -0,0 +1,20 @@
+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 PyTorchHandlerMigrator20231201(base_migrator.BaseModelHandlerMigrator):
+    source_version = "2023-12-01"
+    target_version = "2025-03-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.PyTorchModelBlobOptions, model_blob_metadata.options)
+        model_blob_options["multiple_inputs"] = True
+        model_meta.models[name].options = model_blob_options

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_handlers_migrator/tensorflow_migrator_2025_01_01.py

@@ -0,0 +1,19 @@
+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 TensorflowHandlerMigrator20250101(base_migrator.BaseModelHandlerMigrator):
+    source_version = "2025-01-01"
+    target_version = "2025-03-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)
+        model_blob_options["multiple_inputs"] = True
+        model_meta.models[name].options = model_blob_options

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

@@ -0,0 +1,20 @@
+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 TorchScriptHandlerMigrator20231201(base_migrator.BaseModelHandlerMigrator):
+    source_version = "2023-12-01"
+    target_version = "2025-03-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.PyTorchModelBlobOptions, model_blob_metadata.options)
+        model_blob_options["multiple_inputs"] = True
+        model_meta.models[name].options = model_blob_options

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

@@ -1,2 +1 @@
 REQUIREMENTS = ['cloudpickle>=2.0.0']
-ALL_REQUIREMENTS=['cloudpickle>=2.0.0']

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

@@ -48,6 +48,7 @@ def create_model_metadata(
     ext_modules: Optional[List[ModuleType]] = None,
     conda_dependencies: Optional[List[str]] = None,
     pip_requirements: Optional[List[str]] = None,
+    artifact_repository_map: Optional[Dict[str, str]] = None,
     python_version: Optional[str] = None,
     task: model_types.Task = model_types.Task.UNKNOWN,
     **kwargs: Any,
@@ -67,6 +68,7 @@ def create_model_metadata(
         ext_modules: List of names of modules that need to be pickled with the model. Defaults to None.
         conda_dependencies: List of conda requirements for running the model. Defaults to None.
         pip_requirements: List of pip Python packages requirements for running the model. Defaults to None.
+        artifact_repository_map: A dict mapping from package channel to artifact repository name.
         python_version: A string of python version where model is run. Used for user override. If specified as None,
             current version would be captured. Defaults to None.
         task: The task of the Model Version. It is an enum class Task with values TABULAR_REGRESSION,
@@ -102,6 +104,7 @@ def create_model_metadata(
     env = _create_env_for_model_metadata(
         conda_dependencies=conda_dependencies,
         pip_requirements=pip_requirements,
+        artifact_repository_map=artifact_repository_map,
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
     )
@@ -151,6 +154,7 @@ def _create_env_for_model_metadata(
     *,
     conda_dependencies: Optional[List[str]] = None,
     pip_requirements: Optional[List[str]] = None,
+    artifact_repository_map: Optional[Dict[str, str]] = None,
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
 ) -> model_env.ModelEnv:
@@ -159,6 +163,7 @@ def _create_env_for_model_metadata(
     # Mypy doesn't like getter and setter have different types. See python/mypy #3004
     env.conda_dependencies = conda_dependencies  # type: ignore[assignment]
     env.pip_requirements = pip_requirements  # type: ignore[assignment]
+    env.artifact_repository_map = artifact_repository_map
     env.python_version = python_version  # type: ignore[assignment]
     env.snowpark_ml_version = snowml_env.VERSION
 
@@ -331,7 +336,6 @@ class ModelMetadata:
                 "function_properties": self.function_properties,
             }
         )
-
         with open(model_yaml_path, "w", encoding="utf-8") as out:
             yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
             yaml.safe_dump(model_dict, stream=out, default_flow_style=False)
@@ -352,7 +356,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

@@ -18,6 +18,7 @@ class FunctionProperties(Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
+    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):
@@ -28,6 +29,7 @@ class ModelRuntimeDict(TypedDict):
 class ModelEnvDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
+    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
     python_version: Required[str]
     cuda_version: NotRequired[Optional[str]]
     snowpark_ml_version: Required[str]
@@ -44,6 +46,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):
@@ -58,8 +63,17 @@ class XgboostModelBlobOptions(BaseModelBlobOptions):
     xgb_estimator_type: Required[str]
 
 
+class PyTorchModelBlobOptions(BaseModelBlobOptions):
+    multiple_inputs: Required[bool]
+
+
+class TorchScriptModelBlobOptions(BaseModelBlobOptions):
+    multiple_inputs: Required[bool]
+
+
 class TensorflowModelBlobOptions(BaseModelBlobOptions):
     save_format: Required[str]
+    multiple_inputs: Required[bool]
 
 
 class SentenceTransformersModelBlobOptions(BaseModelBlobOptions):
@@ -71,6 +85,8 @@ ModelBlobOptions = Union[
     HuggingFacePipelineModelBlobOptions,
     MLFlowModelBlobOptions,
     XgboostModelBlobOptions,
+    PyTorchModelBlobOptions,
+    TorchScriptModelBlobOptions,
     TensorflowModelBlobOptions,
     SentenceTransformersModelBlobOptions,
 ]

snowflake/ml/model/_packager/model_packager.py

@@ -43,13 +43,13 @@ class ModelPackager:
         metadata: Optional[Dict[str, str]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        artifact_repository_map: Optional[Dict[str, str]] = None,
         python_version: Optional[str] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
-        options: Optional[model_types.ModelSaveOption] = None,
+        options: model_types.ModelSaveOption,
         task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:
-
         if (signatures is None) and (sample_input_data is None) and not model_handler.is_auto_signature_model(model):
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_ARGUMENT,
@@ -58,9 +58,6 @@ class ModelPackager:
                 ),
             )
 
-        if not options:
-            options = model_types.BaseModelSaveOption()
-
         handler = model_handler.find_handler(model)
         if handler is None:
             raise snowml_exceptions.SnowflakeMLException(
@@ -77,6 +74,7 @@ class ModelPackager:
             ext_modules=ext_modules,
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
+            artifact_repository_map=artifact_repository_map,
             python_version=python_version,
             task=task,
             **options,

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

@@ -1,2 +1 @@
-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']
+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.12.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']

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,13 +89,18 @@ 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)),
             dependencies=model_meta_schema.ModelRuntimeDependenciesDict(
                 conda=env_dict["conda"],
                 pip=env_dict["pip"],
+                artifact_repository_map=env_dict["artifact_repository_map"]
+                if env_dict.get("artifact_repository_map") is not None
+                else {},
             ),
         )
 
@@ -109,6 +115,7 @@ class ModelRuntime:
         env.python_version = meta_env.python_version
         env.cuda_version = meta_env.cuda_version
         env.snowpark_ml_version = meta_env.snowpark_ml_version
+        env.artifact_repository_map = meta_env.artifact_repository_map
 
         conda_env_rel_path = pathlib.PurePosixPath(loaded_dict["dependencies"]["conda"])
         pip_requirements_rel_path = pathlib.PurePosixPath(loaded_dict["dependencies"]["pip"])

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

@@ -24,7 +24,11 @@ def get_task_skl(model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pi
     from sklearn.base import is_classifier, is_regressor
 
     if type_utils.LazyType("sklearn.pipeline.Pipeline").isinstance(model):
-        return type_hints.Task.UNKNOWN
+        if hasattr(model, "predict_proba") or hasattr(model, "predict"):
+            model = model.steps[-1][1]  # type: ignore[attr-defined]
+            return _get_model_task(model)
+        else:
+            return type_hints.Task.UNKNOWN
     if is_regressor(model):
         return type_hints.Task.TABULAR_REGRESSION
     if is_classifier(model):

snowflake/ml/model/_signatures/builtins_handler.py

@@ -14,21 +14,32 @@ from snowflake.ml.model._signatures import base_handler, core, pandas_handler
 
 
 class ListOfBuiltinHandler(base_handler.BaseDataHandler[model_types._SupportedBuiltinsList]):
+    @staticmethod
+    def _can_handle_element(
+        element: model_types._SupportedBuiltins,
+    ) -> TypeGuard[model_types._SupportedBuiltins]:
+        if isinstance(element, abc.Sequence) and not isinstance(element, str):
+            for sub_element in element:
+                if not ListOfBuiltinHandler._can_handle_element(sub_element):
+                    return False
+            return True
+        elif isinstance(element, abc.Mapping):
+            for key, value in element.items():
+                if not isinstance(key, str):
+                    return False
+                if not ListOfBuiltinHandler._can_handle_element(value):
+                    return False
+            return True
+        else:
+            return isinstance(element, (int, float, bool, str, datetime.datetime))
+
     @staticmethod
     def can_handle(data: model_types.SupportedDataType) -> TypeGuard[model_types._SupportedBuiltinsList]:
         if not isinstance(data, abc.Sequence) or isinstance(data, str):
             return False
         if len(data) == 0:
             return False
-        can_handle = True
-        for element in data:
-            # String is a Sequence but we take them as an whole
-            if isinstance(element, abc.Sequence) and not isinstance(element, str):
-                can_handle = ListOfBuiltinHandler.can_handle(element)
-            elif not isinstance(element, (int, float, bool, str, datetime.datetime)):
-                can_handle = False
-                break
-        return can_handle
+        return ListOfBuiltinHandler._can_handle_element(data)
 
     @staticmethod
     def count(data: model_types._SupportedBuiltinsList) -> int: