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

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

@@ -10,10 +10,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,
+    pytorch_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,
@@ -21,7 +25,6 @@ from snowflake.ml.model._signatures import (
 )
 
 if TYPE_CHECKING:
-    import sentence_transformers  # noqa: F401
     import torch
 
 
@@ -33,9 +36,11 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
     """
 
     HANDLER_TYPE = "pytorch"
-    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": pytorch_migrator_2023_12_01.PyTorchHandlerMigrator20231201
+    }
 
     MODEL_BLOB_FILE_OR_DIR = "model.pt"
     DEFAULT_TARGET_METHODS = ["forward"]
@@ -89,22 +94,33 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
                 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 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)
 
-                if isinstance(predictions_df, torch.Tensor):
-                    predictions_df = [predictions_df]
                 return predictions_df
 
             model_meta = handlers_utils.validate_signature(
@@ -127,6 +143,7 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
+            options=model_meta_schema.PyTorchModelBlobOptions(multiple_inputs=multiple_inputs),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -172,6 +189,10 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
             raw_model: "torch.nn.Module",
             model_meta: model_meta_api.ModelMetadata,
         ) -> Type[custom_model.CustomModel]:
+            multiple_inputs = cast(
+                model_meta_schema.PyTorchModelBlobOptions, model_meta.models[model_meta.name].options
+            )["multiple_inputs"]
+
             def fn_factory(
                 raw_model: "torch.nn.Module",
                 signature: model_signature.ModelSignature,
@@ -183,19 +204,28 @@ class PyTorchHandler(_base.BaseModelHandler["torch.nn.Module"]):
                         raise ValueError("Tensor cannot handle null values.")
 
                     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):
+                            st = [element.cuda() for element in st]
 
-                    if kwargs.get("use_gpu", False):
-                        t = [element.cuda() for element in t]
+                        with torch.no_grad():
+                            res = getattr(raw_model, target_method)(*st)
 
-                    with torch.no_grad():
-                        res = getattr(raw_model, target_method)(*t)
+                        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()
 
-                    if isinstance(res, torch.Tensor):
-                        res = [res]
+                        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/sklearn.py

@@ -297,10 +297,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                         df = handlers_utils.convert_explanations_to_2D_df(raw_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
-                            }
+                            dtype_map = {spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in signature.inputs}
 
                             if isinstance(X, pd.DataFrame):
                                 X = X.astype(dtype_map, copy=False)

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

@@ -307,8 +307,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                             except TypeError:
                                 try:
                                     dtype_map = {
-                                        spec.name: spec.as_dtype(force_numpy_dtype=True)  # type: ignore[attr-defined]
-                                        for spec in signature.inputs
+                                        spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in signature.inputs
                                     }
 
                                     if isinstance(X, pd.DataFrame):

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

@@ -1,7 +1,6 @@
 import os
 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
@@ -13,6 +12,7 @@ from snowflake.ml.model._packager.model_handlers import _base, _utils as handler
 from snowflake.ml.model._packager.model_handlers_migrator import (
     base_migrator,
     tensorflow_migrator_2023_12_01,
+    tensorflow_migrator_2025_01_01,
 )
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
@@ -20,7 +20,6 @@ from snowflake.ml.model._packager.model_meta import (
     model_meta_schema,
 )
 from snowflake.ml.model._signatures import (
-    numpy_handler,
     tensorflow_handler,
     utils as model_signature_utils,
 )
@@ -37,10 +36,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
     """
 
     HANDLER_TYPE = "tensorflow"
-    HANDLER_VERSION = "2025-01-01"
-    _MIN_SNOWPARK_ML_VERSION = "1.7.5"
+    HANDLER_VERSION = "2025-03-01"
+    _MIN_SNOWPARK_ML_VERSION = "1.8.0"
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
-        "2023-12-01": tensorflow_migrator_2023_12_01.TensorflowHandlerMigrator20231201
+        "2023-12-01": tensorflow_migrator_2023_12_01.TensorflowHandlerMigrator20231201,
+        "2025-01-01": tensorflow_migrator_2025_01_01.TensorflowHandlerMigrator20250101,
     }
 
     MODEL_BLOB_FILE_OR_DIR = "model"
@@ -112,25 +112,35 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
                 default_target_methods=default_target_methods,
             )
 
+            multiple_inputs = kwargs.get("multiple_inputs", False)
+
             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:
-                if not tensorflow_handler.SeqOfTensorflowTensorHandler.can_handle(sample_input_data):
-                    sample_input_data = tensorflow_handler.SeqOfTensorflowTensorHandler.convert_from_df(
-                        model_signature._convert_local_data_to_df(sample_input_data)
-                    )
+                if multiple_inputs:
+                    if not tensorflow_handler.SeqOfTensorflowTensorHandler.can_handle(sample_input_data):
+                        sample_input_data = tensorflow_handler.SeqOfTensorflowTensorHandler.convert_from_df(
+                            model_signature._convert_local_data_to_df(sample_input_data)
+                        )
+                else:
+                    if not tensorflow_handler.TensorflowTensorHandler.can_handle(sample_input_data):
+                        sample_input_data = tensorflow_handler.TensorflowTensorHandler.convert_from_df(
+                            model_signature._convert_local_data_to_df(sample_input_data)
+                        )
 
                 target_method = getattr(model, target_method_name, None)
                 assert callable(target_method)
                 for tensor in sample_input_data:
                     tensorflow.stop_gradient(tensor)
-                predictions_df = target_method(*sample_input_data)
-
-                if isinstance(predictions_df, (tensorflow.Tensor, tensorflow.Variable, np.ndarray)):
-                    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
 
@@ -159,7 +169,9 @@ 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(save_format=save_format),
+            options=model_meta_schema.TensorflowModelBlobOptions(
+                save_format=save_format, multiple_inputs=multiple_inputs
+            ),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -219,6 +231,10 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             raw_model: "tensorflow.Module",
             model_meta: model_meta_api.ModelMetadata,
         ) -> Type[custom_model.CustomModel]:
+            multiple_inputs = cast(
+                model_meta_schema.TensorflowModelBlobOptions, model_meta.models[model_meta.name].options
+            )["multiple_inputs"]
+
             def fn_factory(
                 raw_model: "tensorflow.Module",
                 signature: model_signature.ModelSignature,
@@ -229,21 +245,25 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
                     if X.isnull().any(axis=None):
                         raise ValueError("Tensor cannot handle null values.")
 
-                    t = tensorflow_handler.SeqOfTensorflowTensorHandler.convert_from_df(X, signature.inputs)
+                    if multiple_inputs:
+                        t = tensorflow_handler.SeqOfTensorflowTensorHandler.convert_from_df(X, signature.inputs)
 
-                    for tensor in t:
-                        tensorflow.stop_gradient(tensor)
-                    res = getattr(raw_model, target_method)(*t)
+                        for tensor in t:
+                            tensorflow.stop_gradient(tensor)
+                        res = getattr(raw_model, target_method)(*t)
 
-                    if isinstance(res, (tensorflow.Tensor, tensorflow.Variable, np.ndarray)):
-                        res = [res]
-
-                    if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
-                        # In case of running on CPU, it will return numpy array
-                        df = numpy_handler.SeqOfNumpyArrayHandler.convert_to_df(res)
+                        if not isinstance(res, tuple):
+                            res = [res]
                     else:
-                        df = tensorflow_handler.SeqOfTensorflowTensorHandler.convert_to_df(res)
-                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+                        t = tensorflow_handler.TensorflowTensorHandler.convert_from_df(X, signature.inputs)
+
+                        tensorflow.stop_gradient(t)
+                        res = getattr(raw_model, target_method)(t)
+
+                    return model_signature_utils.rename_pandas_df(
+                        model_signature._convert_local_data_to_df(res, ensure_serializable=True),
+                        features=signature.outputs,
+                    )
 
                 return fn
 

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_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)

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]
@@ -61,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):
@@ -74,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', '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']
+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']