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

snowflake/ml/model/_signatures/tensorflow_handler.py

@@ -1,5 +1,5 @@
 from collections import abc
-from typing import TYPE_CHECKING, List, Literal, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Literal, Optional, Sequence, Union
 
 import numpy as np
 import pandas as pd
@@ -11,12 +11,62 @@ from snowflake.ml._internal.exceptions import (
     exceptions as snowml_exceptions,
 )
 from snowflake.ml.model import type_hints as model_types
-from snowflake.ml.model._signatures import base_handler, core
+from snowflake.ml.model._signatures import base_handler, core, numpy_handler
 
 if TYPE_CHECKING:
     import tensorflow
 
 
+class TensorflowTensorHandler(base_handler.BaseDataHandler[Union["tensorflow.Tensor", "tensorflow.Variable"]]):
+    @staticmethod
+    def can_handle(
+        data: model_types.SupportedDataType,
+    ) -> TypeGuard[Union["tensorflow.Tensor", "tensorflow.Variable"]]:
+        return type_utils.LazyType("tensorflow.Tensor").isinstance(data) or type_utils.LazyType(
+            "tensorflow.Variable"
+        ).isinstance(data)
+
+    @staticmethod
+    def count(data: Union["tensorflow.Tensor", "tensorflow.Variable"]) -> int:
+        return numpy_handler.NumpyArrayHandler.count(data.numpy())
+
+    @staticmethod
+    def truncate(
+        data: Union["tensorflow.Tensor", "tensorflow.Variable"], length: int
+    ) -> Union["tensorflow.Tensor", "tensorflow.Variable"]:
+        return data[: min(TensorflowTensorHandler.count(data), length)]
+
+    @staticmethod
+    def validate(data: Union["tensorflow.Tensor", "tensorflow.Variable"]) -> None:
+        numpy_handler.NumpyArrayHandler.validate(data.numpy())
+
+    @staticmethod
+    def infer_signature(
+        data: Union["tensorflow.Tensor", "tensorflow.Variable"], role: Literal["input", "output"]
+    ) -> Sequence[core.BaseFeatureSpec]:
+        return numpy_handler.NumpyArrayHandler.infer_signature(data.numpy(), role=role)
+
+    @staticmethod
+    def convert_to_df(
+        data: Union["tensorflow.Tensor", "tensorflow.Variable"], ensure_serializable: bool = True
+    ) -> pd.DataFrame:
+        return numpy_handler.NumpyArrayHandler.convert_to_df(data.numpy(), ensure_serializable=ensure_serializable)
+
+    @staticmethod
+    def convert_from_df(
+        df: pd.DataFrame, features: Optional[Sequence[core.BaseFeatureSpec]] = None
+    ) -> Union["tensorflow.Tensor", "tensorflow.Variable"]:
+        import tensorflow as tf
+
+        if features is None:
+            if any(dtype == np.dtype("O") for dtype in df.dtypes):
+                return tf.convert_to_tensor(np.array(df.to_numpy().tolist()))
+            return tf.convert_to_tensor(df.to_numpy())
+
+        assert isinstance(features[0], core.FeatureSpec)
+        return tf.convert_to_tensor(np.array(df.to_numpy().tolist()), dtype=features[0]._dtype._numpy_type)
+
+
 class SeqOfTensorflowTensorHandler(
     base_handler.BaseDataHandler[Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]]]
 ):
@@ -28,35 +78,12 @@ class SeqOfTensorflowTensorHandler(
             return False
         if len(data) == 0:
             return False
-        if type_utils.LazyType("tensorflow.Tensor").isinstance(data[0]) or type_utils.LazyType(
-            "tensorflow.Variable"
-        ).isinstance(data[0]):
-            return all(
-                type_utils.LazyType("tensorflow.Tensor").isinstance(data_col)
-                or type_utils.LazyType("tensorflow.Variable").isinstance(data_col)
-                for data_col in data
-            )
-        return False
+
+        return all(TensorflowTensorHandler.can_handle(data_col) for data_col in data)
 
     @staticmethod
     def count(data: Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]]) -> int:
-        import tensorflow as tf
-
-        rows = []
-        for data_col in data:
-            shapes = data_col.shape.as_list()
-            if data_col.shape == tf.TensorShape(None) or (not shapes) or (shapes[0] is None):
-                # Unknown shape array
-                raise snowml_exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_DATA,
-                    original_exception=ValueError("Data Validation Error: Unknown shape data is found."),
-                )
-            # Make mypy happy
-            assert isinstance(shapes[0], int)
-
-            rows.append(shapes[0])
-
-        return min(rows)
+        return min(TensorflowTensorHandler.count(data_col) for data_col in data)
 
     @staticmethod
     def truncate(
@@ -66,49 +93,14 @@ class SeqOfTensorflowTensorHandler(
 
     @staticmethod
     def validate(data: Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]]) -> None:
-        import tensorflow as tf
-
         for data_col in data:
-            if data_col.shape == tf.TensorShape(None) or any(dim is None for dim in data_col.shape.as_list()):
-                # Unknown shape array
-                raise snowml_exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_DATA,
-                    original_exception=ValueError("Data Validation Error: Unknown shape data is found."),
-                )
-
-            if data_col.shape == tf.TensorShape([0]):
-                # Empty array
-                raise snowml_exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_DATA,
-                    original_exception=ValueError("Data Validation Error: Empty data is found."),
-                )
-
-            if data_col.shape == tf.TensorShape([1]) or data_col.shape == tf.TensorShape([]):
-                # scalar
-                raise snowml_exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_DATA,
-                    original_exception=ValueError("Data Validation Error: Scalar data is found."),
-                )
+            TensorflowTensorHandler.validate(data_col)
 
     @staticmethod
     def infer_signature(
         data: Sequence[Union["tensorflow.Tensor", "tensorflow.Variable"]], role: Literal["input", "output"]
     ) -> Sequence[core.BaseFeatureSpec]:
-        feature_prefix = f"{SeqOfTensorflowTensorHandler.FEATURE_PREFIX}_"
-        features: List[core.BaseFeatureSpec] = []
-        role_prefix = (
-            SeqOfTensorflowTensorHandler.INPUT_PREFIX if role == "input" else SeqOfTensorflowTensorHandler.OUTPUT_PREFIX
-        ) + "_"
-
-        for i, data_col in enumerate(data):
-            dtype = core.DataType.from_numpy_type(data_col.dtype.as_numpy_dtype)
-            ft_name = f"{role_prefix}{feature_prefix}{i}"
-            if len(data_col.shape) == 1:
-                features.append(core.FeatureSpec(dtype=dtype, name=ft_name, nullable=False))
-            else:
-                ft_shape = tuple(data_col.shape[1:])
-                features.append(core.FeatureSpec(dtype=dtype, name=ft_name, shape=ft_shape, nullable=False))
-        return features
+        return numpy_handler.SeqOfNumpyArrayHandler.infer_signature([data_col.numpy() for data_col in data], role=role)
 
     @staticmethod
     def convert_to_df(
@@ -129,8 +121,10 @@ class SeqOfTensorflowTensorHandler(
             for feature in features:
                 if isinstance(feature, core.FeatureGroupSpec):
                     raise snowml_exceptions.SnowflakeMLException(
-                        error_code=error_codes.NOT_IMPLEMENTED,
-                        original_exception=NotImplementedError("FeatureGroupSpec is not supported."),
+                        error_code=error_codes.INVALID_DATA_TYPE,
+                        original_exception=NotImplementedError(
+                            "FeatureGroupSpec is not supported when converting to Tensorflow tensor."
+                        ),
                     )
                 assert isinstance(feature, core.FeatureSpec), "Invalid feature kind."
                 res.append(

snowflake/ml/model/_signatures/utils.py

@@ -135,7 +135,16 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
                 core.FeatureSpec(name="inputs", dtype=core.DataType.STRING),
             ],
             outputs=[
-                core.FeatureSpec(name="outputs", dtype=core.DataType.STRING),
+                core.FeatureGroupSpec(
+                    name="outputs",
+                    specs=[
+                        core.FeatureSpec(name="sequence", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="score", dtype=core.DataType.DOUBLE),
+                        core.FeatureSpec(name="token", dtype=core.DataType.INT64),
+                        core.FeatureSpec(name="token_str", dtype=core.DataType.STRING),
+                    ],
+                    shape=(-1,),
+                ),
             ],
         )
 
@@ -144,7 +153,18 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
         return core.ModelSignature(
             inputs=[core.FeatureSpec(name="inputs", dtype=core.DataType.STRING)],
             outputs=[
-                core.FeatureSpec(name="outputs", dtype=core.DataType.STRING),
+                core.FeatureGroupSpec(
+                    name="outputs",
+                    specs=[
+                        core.FeatureSpec(name="word", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="score", dtype=core.DataType.DOUBLE),
+                        core.FeatureSpec(name="entity", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="index", dtype=core.DataType.INT64),
+                        core.FeatureSpec(name="start", dtype=core.DataType.INT64),
+                        core.FeatureSpec(name="end", dtype=core.DataType.INT64),
+                    ],
+                    shape=(-1,),
+                ),
             ],
         )
 
@@ -171,7 +191,16 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
                 core.FeatureSpec(name="context", dtype=core.DataType.STRING),
             ],
             outputs=[
-                core.FeatureSpec(name="outputs", dtype=core.DataType.STRING),
+                core.FeatureGroupSpec(
+                    name="answers",
+                    specs=[
+                        core.FeatureSpec(name="score", dtype=core.DataType.DOUBLE),
+                        core.FeatureSpec(name="start", dtype=core.DataType.INT64),
+                        core.FeatureSpec(name="end", dtype=core.DataType.INT64),
+                        core.FeatureSpec(name="answer", dtype=core.DataType.STRING),
+                    ],
+                    shape=(-1,),
+                ),
             ],
         )
 
@@ -216,17 +245,22 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
             return core.ModelSignature(
                 inputs=[
                     core.FeatureSpec(name="text", dtype=core.DataType.STRING),
-                    core.FeatureSpec(name="text_pair", dtype=core.DataType.STRING),
                 ],
                 outputs=[
-                    core.FeatureSpec(name="outputs", dtype=core.DataType.STRING),
+                    core.FeatureGroupSpec(
+                        name="labels",
+                        specs=[
+                            core.FeatureSpec(name="label", dtype=core.DataType.STRING),
+                            core.FeatureSpec(name="score", dtype=core.DataType.DOUBLE),
+                        ],
+                        shape=(-1,),
+                    ),
                 ],
             )
         # Else, return a dict per input
         return core.ModelSignature(
             inputs=[
                 core.FeatureSpec(name="text", dtype=core.DataType.STRING),
-                core.FeatureSpec(name="text_pair", dtype=core.DataType.STRING),
             ],
             outputs=[
                 core.FeatureSpec(name="label", dtype=core.DataType.STRING),
@@ -243,9 +277,24 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
             )
         # Always generate a list of dict per input
         return core.ModelSignature(
-            inputs=[core.FeatureSpec(name="inputs", dtype=core.DataType.STRING)],
+            inputs=[
+                core.FeatureGroupSpec(
+                    name="inputs",
+                    specs=[
+                        core.FeatureSpec(name="role", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="content", dtype=core.DataType.STRING),
+                    ],
+                    shape=(-1,),
+                ),
+            ],
             outputs=[
-                core.FeatureSpec(name="outputs", dtype=core.DataType.STRING),
+                core.FeatureGroupSpec(
+                    name="outputs",
+                    specs=[
+                        core.FeatureSpec(name="generated_text", dtype=core.DataType.STRING),
+                    ],
+                    shape=(-1,),
+                )
             ],
         )
 
@@ -300,3 +349,66 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any])
 
 def series_dropna(series: pd.Series) -> pd.Series:
     return series.dropna(inplace=False).reset_index(drop=True).convert_dtypes()
+
+
+def infer_list(name: str, data: List[Any]) -> core.BaseFeatureSpec:
+    """Infer the feature specification from a list.
+
+    Args:
+        name: Feature name.
+        data: A list.
+
+    Raises:
+        SnowflakeMLException: ValueError: Raised when empty list is provided.
+
+    Returns:
+        A feature specification.
+    """
+    if not data:
+        raise snowml_exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_DATA,
+            original_exception=ValueError("Data Validation Error: Empty list is found."),
+        )
+
+    if all(isinstance(value, dict) for value in data):
+        ft = infer_dict(name, data[0])
+        ft._name = name
+        ft._shape = (-1,)
+        return ft
+
+    arr = convert_list_to_ndarray(data)
+    arr_dtype = core.DataType.from_numpy_type(arr.dtype)
+
+    return core.FeatureSpec(name=name, dtype=arr_dtype, shape=arr.shape)
+
+
+def infer_dict(name: str, data: Dict[str, Any]) -> core.FeatureGroupSpec:
+    """Infer the feature specification from a dictionary.
+
+    Args:
+        name: Feature name.
+        data: A dictionary.
+
+    Raises:
+        SnowflakeMLException: ValueError: Raised when empty dictionary is provided.
+        SnowflakeMLException: ValueError: Raised when empty list is found in the dictionary.
+
+    Returns:
+        A feature group specification.
+    """
+    if not data:
+        raise snowml_exceptions.SnowflakeMLException(
+            error_code=error_codes.INVALID_DATA,
+            original_exception=ValueError("Data Validation Error: Empty dictionary is found."),
+        )
+
+    specs = []
+    for key, value in data.items():
+        if isinstance(value, list):
+            specs.append(infer_list(key, value))
+        elif isinstance(value, dict):
+            specs.append(infer_dict(key, value))
+        else:
+            specs.append(core.FeatureSpec(name=key, dtype=core.DataType.from_numpy_type(np.array(value).dtype)))
+
+    return core.FeatureGroupSpec(name=name, specs=specs)

snowflake/ml/model/custom_model.py

@@ -76,7 +76,7 @@ class ModelRef:
     def __getattr__(self, method_name: str) -> Any:
         if hasattr(self._model, method_name):
             return MethodRef(self, method_name)
-        raise TypeError(f"Model is does not have {method_name}.")
+        raise AttributeError(f"Method {method_name} not found in model {self._name}.")
 
     def __getstate__(self) -> Dict[str, Any]:
         state = self.__dict__.copy()
@@ -94,7 +94,16 @@ class ModelRef:
 
 class ModelContext:
     """
-    Context for a custom model showing paths to artifacts and mapping between model name and object reference.
+    Context for a custom model storing paths to file artifacts and model object references.
+
+    Keyword argument values can be string file paths or supported in-memory models. Paths and model references
+    can be accessed with dictionary access methods in the custom model.
+
+    For example, in a custom model with `context=ModelContext(my_file='my_file.pkl', my_model=my_model)`,
+    the filepath and model reference can be accessed with `self.context['my_file']` and `self.context['my_model']`
+    in the inference and init methods.
+
+    The use of `artifacts` and `model_refs` arguments is deprecated. Set keyword arguments directly instead.
 
     Attributes:
         artifacts: A dictionary mapping the name of the artifact to its path.
@@ -267,14 +276,14 @@ def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.D
 
 
 def inference_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]
+    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
 ) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     return func
 
 
 def partitioned_inference_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]
+    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
 ) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     func.__dict__["_is_partitioned_inference_api"] = True

snowflake/ml/model/model_signature.py

@@ -32,6 +32,7 @@ from snowflake.ml.model._signatures import (
     base_handler,
     builtins_handler as builtins_handler,
     core,
+    dmatrix_handler,
     numpy_handler,
     pandas_handler,
     pytorch_handler,
@@ -52,8 +53,11 @@ _LOCAL_DATA_HANDLERS: List[Type[base_handler.BaseDataHandler[Any]]] = [
     numpy_handler.NumpyArrayHandler,
     builtins_handler.ListOfBuiltinHandler,
     numpy_handler.SeqOfNumpyArrayHandler,
+    pytorch_handler.PyTorchTensorHandler,
     pytorch_handler.SeqOfPyTorchTensorHandler,
+    tensorflow_handler.TensorflowTensorHandler,
     tensorflow_handler.SeqOfTensorflowTensorHandler,
+    dmatrix_handler.XGBoostDMatrixHandler,
 ]
 _ALL_DATA_HANDLERS = _LOCAL_DATA_HANDLERS + [snowpark_handler.SnowparkDataFrameHandler]
 
@@ -218,7 +222,6 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
         strict: Enable strict validation, this includes value range based validation
 
     Raises:
-        SnowflakeMLException: NotImplementedError: FeatureGroupSpec is not supported.
         SnowflakeMLException: ValueError: Raised when a feature cannot be found.
         SnowflakeMLException: ValueError: Raised when feature is scalar but confront list element.
         SnowflakeMLException: ValueError: Raised when feature type is not aligned in list element.
@@ -236,7 +239,10 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
         except KeyError:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_DATA,
-                original_exception=ValueError(f"Data Validation Error: feature {ft_name} does not exist in data."),
+                original_exception=ValueError(
+                    f"Data Validation Error: feature {ft_name} does not exist in data. "
+                    f"Available columns are {data.columns}."
+                ),
             )
 
         if data_col.isnull().any():
@@ -244,10 +250,15 @@ def _validate_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureS
         df_col_dtype = data_col.dtype
 
         if isinstance(feature, core.FeatureGroupSpec):
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.NOT_IMPLEMENTED,
-                original_exception=NotImplementedError("FeatureGroupSpec is not supported."),
-            )
+            if df_col_dtype != np.dtype("O"):
+                raise snowml_exceptions.SnowflakeMLException(
+                    error_code=error_codes.INVALID_DATA,
+                    original_exception=ValueError(
+                        f"Data Validation Error in feature group {ft_name}: "
+                        + f"It needs to be a dictionary or list of dictionary, but get {df_col_dtype}."
+                    ),
+                )
+            continue
 
         assert isinstance(feature, core.FeatureSpec)  # assert for mypy.
         ft_type = feature._dtype
@@ -437,7 +448,6 @@ def _validate_snowpark_data(
         strict: Enable strict validation, this includes value range based validation.
 
     Raises:
-        SnowflakeMLException: NotImplementedError: FeatureGroupSpec is not supported.
         SnowflakeMLException: ValueError: Raised when confronting invalid feature.
         SnowflakeMLException: ValueError: Raised when a feature cannot be found.
 
@@ -467,10 +477,15 @@ def _validate_snowpark_data(
                 if field.name == ft_name:
                     found = True
                     if isinstance(feature, core.FeatureGroupSpec):
-                        raise snowml_exceptions.SnowflakeMLException(
-                            error_code=error_codes.NOT_IMPLEMENTED,
-                            original_exception=NotImplementedError("FeatureGroupSpec is not supported."),
-                        )
+                        if not isinstance(field.datatype, (spt.ArrayType, spt.StructType, spt.VariantType)):
+                            errors[identifier_rule].append(
+                                ValueError(
+                                    f"Data Validation Error in feature group {feature.name}: "
+                                    + f"Feature expects {feature.as_snowpark_type()},"
+                                    + f" while {field.name} has type {field.datatype}."
+                                ),
+                            )
+                        continue
                     assert isinstance(feature, core.FeatureSpec)  # mypy
                     ft_type = feature._dtype
                     field_data_type = field.datatype
@@ -644,11 +659,14 @@ def _validate_snowpark_type_feature(
             )
 
 
-def _convert_local_data_to_df(data: model_types.SupportedLocalDataType) -> pd.DataFrame:
+def _convert_local_data_to_df(
+    data: model_types.SupportedLocalDataType, ensure_serializable: bool = False
+) -> pd.DataFrame:
     """Convert local data to pandas DataFrame or Snowpark DataFrame
 
     Args:
         data: The provided data.
+        ensure_serializable: Ensure the data is serializable. Defaults to False.
 
     Raises:
         SnowflakeMLException: NotImplementedError: Raised when data cannot be handled by any data handler.
@@ -660,7 +678,7 @@ def _convert_local_data_to_df(data: model_types.SupportedLocalDataType) -> pd.Da
     for handler in _LOCAL_DATA_HANDLERS:
         if handler.can_handle(data):
             handler.validate(data)
-            df = handler.convert_to_df(data, ensure_serializable=False)
+            df = handler.convert_to_df(data, ensure_serializable=ensure_serializable)
             break
     if df is None:
         raise snowml_exceptions.SnowflakeMLException(

snowflake/ml/model/type_hints.py

@@ -26,7 +26,15 @@ if TYPE_CHECKING:
     from snowflake.ml.modeling.framework import base  # noqa: F401
 
 
-_SupportedBuiltins = Union[int, float, bool, str, bytes, "_SupportedBuiltinsList"]
+_SupportedBuiltins = Union[
+    int,
+    float,
+    bool,
+    str,
+    bytes,
+    Dict[str, Union["_SupportedBuiltins", "_SupportedBuiltinsList"]],
+    "_SupportedBuiltinsList",
+]
 _SupportedNumpyDtype = Union[
     "np.int8",
     "np.int16",
@@ -48,7 +56,7 @@ _SupportedBuiltinsList = Sequence[_SupportedBuiltins]
 _SupportedArrayLike = Union[_SupportedNumpyArray, "torch.Tensor", "tensorflow.Tensor", "tensorflow.Variable"]
 
 SupportedLocalDataType = Union[
-    "pd.DataFrame", _SupportedNumpyArray, Sequence[_SupportedArrayLike], _SupportedBuiltinsList
+    "pd.DataFrame", _SupportedArrayLike, Sequence[_SupportedArrayLike], _SupportedBuiltinsList
 ]
 
 SupportedDataType = Union[SupportedLocalDataType, "snowflake.snowpark.DataFrame"]
@@ -177,16 +185,19 @@ class SNOWModelSaveOptions(BaseModelSaveOption):
 class PyTorchSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    multiple_inputs: NotRequired[bool]
 
 
 class TorchScriptSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    multiple_inputs: NotRequired[bool]
 
 
 class TensorflowSaveOptions(BaseModelSaveOption):
     target_methods: NotRequired[Sequence[str]]
     cuda_version: NotRequired[str]
+    multiple_inputs: NotRequired[bool]
 
 
 class MLFlowSaveOptions(BaseModelSaveOption):

snowflake/ml/modeling/metrics/ranking.py

@@ -98,6 +98,7 @@ def precision_recall_curve(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -245,6 +246,7 @@ def roc_auc_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -348,6 +350,7 @@ def roc_curve(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,

snowflake/ml/modeling/metrics/regression.py

@@ -83,6 +83,7 @@ def d2_absolute_error_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -180,6 +181,7 @@ def d2_pinball_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,
@@ -295,6 +297,7 @@ def explained_variance_score(
         packages=[
             f"cloudpickle=={cloudpickle.__version__}",
             f"scikit-learn=={sklearn_release[0]}.{sklearn_release[1]}.*",
+            f"numpy=={np.__version__}",
             "snowflake-snowpark-python",
         ],
         statement_params=statement_params,

snowflake/ml/modeling/preprocessing/k_bins_discretizer.py

@@ -341,7 +341,7 @@ class KBinsDiscretizer(base.BaseTransformer):
             is_permanent=False,
             name=udf_name,
             replace=True,
-            packages=["numpy"],
+            packages=[f"numpy=={np.__version__}"],
             session=dataset._session,
             statement_params=telemetry.get_statement_params(base.PROJECT, base.SUBPROJECT, self.__class__.__name__),
         )