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

snowflake/ml/model/_signatures/pytorch_handler.py

@@ -1,5 +1,5 @@
 from collections import abc
-from typing import TYPE_CHECKING, List, Literal, Optional, Sequence
+from typing import TYPE_CHECKING, Literal, Optional, Sequence
 
 import numpy as np
 import pandas as pd
@@ -11,12 +11,54 @@ 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 torch
 
 
+class PyTorchTensorHandler(base_handler.BaseDataHandler["torch.Tensor"]):
+    @staticmethod
+    def can_handle(data: model_types.SupportedDataType) -> TypeGuard["torch.Tensor"]:
+        return type_utils.LazyType("torch.Tensor").isinstance(data)
+
+    @staticmethod
+    def count(data: "torch.Tensor") -> int:
+        return data.shape[0]
+
+    @staticmethod
+    def truncate(data: "torch.Tensor", length: int) -> "torch.Tensor":
+        return data[: min(PyTorchTensorHandler.count(data), length)]
+
+    @staticmethod
+    def validate(data: "torch.Tensor") -> None:
+        return numpy_handler.NumpyArrayHandler.validate(data.detach().cpu().numpy())
+
+    @staticmethod
+    def infer_signature(data: "torch.Tensor", role: Literal["input", "output"]) -> Sequence[core.BaseFeatureSpec]:
+        return numpy_handler.NumpyArrayHandler.infer_signature(data.detach().cpu().numpy(), role=role)
+
+    @staticmethod
+    def convert_to_df(data: "torch.Tensor", ensure_serializable: bool = True) -> pd.DataFrame:
+        return numpy_handler.NumpyArrayHandler.convert_to_df(
+            data.detach().cpu().numpy(), ensure_serializable=ensure_serializable
+        )
+
+    @staticmethod
+    def convert_from_df(df: pd.DataFrame, features: Optional[Sequence[core.BaseFeatureSpec]] = None) -> "torch.Tensor":
+        import torch
+
+        if features is None:
+            if any(dtype == np.dtype("O") for dtype in df.dtypes):
+                return torch.from_numpy(np.array(df.to_numpy().tolist()))
+            return torch.from_numpy(df.to_numpy())
+
+        assert isinstance(features[0], core.FeatureSpec)
+        return torch.from_numpy(
+            np.array(df.to_numpy().tolist(), dtype=features[0]._dtype._numpy_type),
+        )
+
+
 class SeqOfPyTorchTensorHandler(base_handler.BaseDataHandler[Sequence["torch.Tensor"]]):
     @staticmethod
     def can_handle(data: model_types.SupportedDataType) -> TypeGuard[Sequence["torch.Tensor"]]:
@@ -24,56 +66,28 @@ class SeqOfPyTorchTensorHandler(base_handler.BaseDataHandler[Sequence["torch.Ten
             return False
         if len(data) == 0:
             return False
-        if type_utils.LazyType("torch.Tensor").isinstance(data[0]):
-            return all(type_utils.LazyType("torch.Tensor").isinstance(data_col) for data_col in data)
-        return False
+        return all(PyTorchTensorHandler.can_handle(data_col) for data_col in data)
 
     @staticmethod
     def count(data: Sequence["torch.Tensor"]) -> int:
-        return min(data_col.shape[0] for data_col in data)
+        return min(PyTorchTensorHandler.count(data_col) for data_col in data)
 
     @staticmethod
     def truncate(data: Sequence["torch.Tensor"], length: int) -> Sequence["torch.Tensor"]:
-        return [data_col[: min(SeqOfPyTorchTensorHandler.count(data), 10)] for data_col in data]
+        return [data_col[: min(SeqOfPyTorchTensorHandler.count(data), length)] for data_col in data]
 
     @staticmethod
     def validate(data: Sequence["torch.Tensor"]) -> None:
-        import torch
-
         for data_col in data:
-            if data_col.shape == torch.Size([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 == torch.Size([1]):
-                # scalar
-                raise snowml_exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_DATA,
-                    original_exception=ValueError("Data Validation Error: Scalar data is found."),
-                )
+            PyTorchTensorHandler.validate(data_col)
 
     @staticmethod
     def infer_signature(
         data: Sequence["torch.Tensor"], role: Literal["input", "output"]
     ) -> Sequence[core.BaseFeatureSpec]:
-        feature_prefix = f"{SeqOfPyTorchTensorHandler.FEATURE_PREFIX}_"
-        features: List[core.BaseFeatureSpec] = []
-        role_prefix = (
-            SeqOfPyTorchTensorHandler.INPUT_PREFIX if role == "input" else SeqOfPyTorchTensorHandler.OUTPUT_PREFIX
-        ) + "_"
-
-        for i, data_col in enumerate(data):
-            dtype = core.DataType.from_torch_type(data_col.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.detach().cpu().numpy() for data_col in data], role=role
+        )
 
     @staticmethod
     def convert_to_df(data: Sequence["torch.Tensor"], ensure_serializable: bool = True) -> pd.DataFrame:
@@ -81,8 +95,8 @@ class SeqOfPyTorchTensorHandler(base_handler.BaseDataHandler[Sequence["torch.Ten
         # the content is still numpy array so that the type could be preserved.
         # But that would not serializable and cannot use as UDF input and output.
         if ensure_serializable:
-            return pd.DataFrame({i: data_col.detach().to("cpu").numpy().tolist() for i, data_col in enumerate(data)})
-        return pd.DataFrame({i: list(data_col.detach().to("cpu").numpy()) for i, data_col in enumerate(data)})
+            return pd.DataFrame({i: data_col.detach().cpu().numpy().tolist() for i, data_col in enumerate(data)})
+        return pd.DataFrame({i: list(data_col.detach().cpu().numpy()) for i, data_col in enumerate(data)})
 
     @staticmethod
     def convert_from_df(
@@ -95,8 +109,10 @@ class SeqOfPyTorchTensorHandler(base_handler.BaseDataHandler[Sequence["torch.Ten
             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(torch.from_numpy(np.stack(df[feature.name].to_numpy()).astype(feature._dtype._numpy_type)))

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -65,12 +65,6 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         dtype_map = {}
         if features:
             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."),
-                    )
-                assert isinstance(feature, core.FeatureSpec), "Invalid feature kind."
                 dtype_map[feature.name] = feature.as_dtype()
         df_local = data.to_pandas()
 
@@ -122,12 +116,6 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         column_names = []
         columns = []
         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."),
-                )
-            assert isinstance(feature, core.FeatureSpec), "Invalid feature kind."
             column_names.append(identifier.get_inferred_name(feature.name))
             columns.append(F.col(identifier.get_inferred_name(feature.name)).cast(feature.as_snowpark_type()))
 

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