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

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

@@ -30,10 +30,7 @@ from snowflake.ml.model._packager.model_meta import (
     model_meta as model_meta_api,
     model_meta_schema,
 )
-from snowflake.ml.model._signatures import (
-    builtins_handler,
-    utils as model_signature_utils,
-)
+from snowflake.ml.model._signatures import utils as model_signature_utils
 from snowflake.ml.model.models import huggingface_pipeline
 from snowflake.snowpark._internal import utils as snowpark_utils
 
@@ -66,16 +63,16 @@ def get_requirements_from_task(task: str, spcs_only: bool = False) -> List[model
     return []
 
 
-class NumpyEncoder(json.JSONEncoder):
-    # This is a JSON encoder class to ensure the output from Huggingface pipeline is JSON serializable.
-    # What it covers is numpy object.
-    def default(self, z: object) -> object:
-        if isinstance(z, np.number):
-            if np.can_cast(z, np.int64, casting="safe"):
-                return int(z)
-            elif np.can_cast(z, np.float64, casting="safe"):
-                return z.astype(np.float64)
-        return super().default(z)
+def sanitize_output(data: Any) -> Any:
+    if isinstance(data, np.number):
+        return data.item()
+    if isinstance(data, np.ndarray):
+        return sanitize_output(data.tolist())
+    if isinstance(data, list):
+        return [sanitize_output(x) for x in data]
+    if isinstance(data, dict):
+        return {k: sanitize_output(v) for k, v in data.items()}
+    return data
 
 
 @final
@@ -410,13 +407,17 @@ class HuggingFacePipelineHandler(
                                 )
                                 for conv_data in X.to_dict("records")
                             ]
-                        elif len(signature.inputs) == 1:
-                            input_data = X.to_dict("list")[signature.inputs[0].name]
                         else:
                             if isinstance(raw_model, transformers.TableQuestionAnsweringPipeline):
                                 X["table"] = X["table"].apply(json.loads)
 
-                            input_data = X.to_dict("records")
+                            # Most pipelines if it is expecting more than one arguments,
+                            # it is expecting a list of dict, where each dict has keys corresponding to the argument.
+                            if len(signature.inputs) > 1:
+                                input_data = X.to_dict("records")
+                            # If it is only expecting one argument, Then it is expecting a list of something.
+                            else:
+                                input_data = X[signature.inputs[0].name].to_list()
                         temp_res = getattr(raw_model, target_method)(input_data)
 
                     # Some huggingface pipeline will omit the outer list when there is only 1 input.
@@ -439,7 +440,6 @@ class HuggingFacePipelineHandler(
                                 ),
                             )
                             and X.shape[0] == 1
-                            and isinstance(temp_res[0], dict)
                         )
                     ):
                         temp_res = [temp_res]
@@ -453,14 +453,18 @@ class HuggingFacePipelineHandler(
                         temp_res = [[conv.generated_responses] for conv in temp_res]
 
                     # To concat those who outputs a list with one input.
-                    if builtins_handler.ListOfBuiltinHandler.can_handle(temp_res):
-                        res = builtins_handler.ListOfBuiltinHandler.convert_to_df(temp_res)
-                    elif isinstance(temp_res[0], dict):
+                    if isinstance(temp_res[0], list):
+                        if isinstance(temp_res[0][0], dict):
+                            res = pd.DataFrame({0: temp_res})
+                        else:
+                            res = pd.DataFrame(temp_res)
+                    else:
                         res = pd.DataFrame(temp_res)
-                    elif isinstance(temp_res[0], list):
-                        res = pd.DataFrame([json.dumps(output, cls=NumpyEncoder) for output in temp_res])
+
+                    if hasattr(res, "map"):
+                        res = res.map(sanitize_output)
                     else:
-                        raise ValueError(f"Cannot parse output {temp_res} from pipeline object")
+                        res = res.applymap(sanitize_output)
 
                     return model_signature_utils.rename_pandas_df(data=res, features=signature.outputs)
 

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

@@ -191,11 +191,7 @@ class KerasHandler(_base.BaseModelHandler["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)
-                }
+                dtype_map = {spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in signature.inputs}
 
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:

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

@@ -57,6 +57,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         "predict_proba",
         "predict_log_proba",
         "decision_function",
+        "score_samples",
     ]
     EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
 
@@ -74,10 +75,6 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             and (
                 not type_utils.LazyType("lightgbm.LGBMModel").isinstance(model)
             )  # LGBMModel is actually a BaseEstimator
-            and any(
-                (hasattr(model, method) and callable(getattr(model, method, None)))
-                for method in cls.DEFAULT_TARGET_METHODS
-            )
         )
 
     @classmethod
@@ -297,10 +294,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']
+REQUIREMENTS = ['cloudpickle>=2.0.0,<3']