snowflake-ml-python 1.7.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

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

@@ -19,12 +19,26 @@ from snowflake.ml.model._packager.model_meta import (
 )
 from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
+from snowflake.ml.modeling._internal.constants import IN_ML_RUNTIME_ENV_VAR
 
 if TYPE_CHECKING:
     import sklearn.base
     import sklearn.pipeline
 
 
+def _unpack_container_runtime_pipeline(model: "sklearn.pipeline.Pipeline") -> "sklearn.pipeline.Pipeline":
+    new_steps = []
+    for step_name, step in model.steps:
+        new_reg = step
+        if hasattr(step, "_sklearn_estimator") and step._sklearn_estimator is not None:
+            # Unpack estimator to open source.
+            new_reg = step._sklearn_estimator
+        new_steps.append((step_name, new_reg))
+
+    model.steps = new_steps
+    return model
+
+
 @final
 class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]]):
     """Handler for scikit-learn based model.
@@ -101,6 +115,10 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             if sample_input_data is None:
                 raise ValueError("Sample input data is required to enable explainability.")
 
+        # If this is a pipeline and we are in the container runtime, check for distributed estimator.
+        if os.getenv(IN_ML_RUNTIME_ENV_VAR) and isinstance(model, sklearn.pipeline.Pipeline):
+            model = _unpack_container_runtime_pipeline(model)
+
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
                 model=model,
@@ -135,7 +153,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             )
 
             model_task_and_output_type = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = model_task_and_output_type.task
+            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output_type.task)
 
             # if users did not ask then we enable if we have background data
             if enable_explainability is None:
@@ -177,6 +195,35 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
+        # if model instance is a pipeline, check the pipeline steps
+        if isinstance(model, sklearn.pipeline.Pipeline):
+            for _, pipeline_step in model.steps:
+                if type_utils.LazyType("lightgbm.LGBMModel").isinstance(pipeline_step) or type_utils.LazyType(
+                    "lightgbm.Booster"
+                ).isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="lightgbm", pip_name="lightgbm"),
+                        ],
+                        check_local_version=True,
+                    )
+                elif type_utils.LazyType("xgboost.XGBModel").isinstance(pipeline_step) or type_utils.LazyType(
+                    "xgboost.Booster"
+                ).isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
+                        ],
+                        check_local_version=True,
+                    )
+                elif type_utils.LazyType("catboost.CatBoost").isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="catboost", pip_name="catboost"),
+                        ],
+                        check_local_version=True,
+                    )
+
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP

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

@@ -138,7 +138,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 enable_explainability = False
             else:
                 model_task_and_output_type = model_task_utils.get_model_task_and_output_type(python_base_obj)
-                model_meta.task = model_task_and_output_type.task
+                model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output_type.task)
                 explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,

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

@@ -13,6 +13,7 @@ from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 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 (
     numpy_handler,
@@ -76,7 +77,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
 
         assert isinstance(model, tensorflow.Module)
 
-        if isinstance(model, tensorflow.keras.Model):
+        is_keras_model = type_utils.LazyType("tensorflow.keras.Model").isinstance(model) or type_utils.LazyType(
+            "tf_keras.Model"
+        ).isinstance(model)
+
+        if is_keras_model:
             default_target_methods = ["predict"]
         else:
             default_target_methods = cls.DEFAULT_TARGET_METHODS
@@ -117,8 +122,14 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
-        if isinstance(model, tensorflow.keras.Model):
+        if is_keras_model:
             tensorflow.keras.models.save_model(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
+            model_meta.env.include_if_absent(
+                [
+                    model_env.ModelDependency(requirement="keras<3", pip_name="keras"),
+                ],
+                check_local_version=False,
+            )
         else:
             tensorflow.saved_model.save(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
 
@@ -127,12 +138,16 @@ 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(is_keras_model=is_keras_model),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
         model_meta.env.include_if_absent(
-            [model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow")], check_local_version=True
+            [
+                model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"),
+            ],
+            check_local_version=True,
         )
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
@@ -150,9 +165,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blobs_metadata = model_meta.models
         model_blob_metadata = model_blobs_metadata[name]
         model_blob_filename = model_blob_metadata.path
-        m = tensorflow.keras.models.load_model(os.path.join(model_blob_path, model_blob_filename), compile=False)
-        if isinstance(m, tensorflow.keras.Model):
-            return m
+        model_blob_options = cast(model_meta_schema.TensorflowModelBlobOptions, model_blob_metadata.options)
+        if model_blob_options.get("is_keras_model", False):
+            m = tensorflow.keras.models.load_model(os.path.join(model_blob_path, model_blob_filename), compile=False)
+        else:
+            m = tensorflow.saved_model.load(os.path.join(model_blob_path, model_blob_filename))
         return cast(tensorflow.Module, m)
 
     @classmethod

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

@@ -23,7 +23,7 @@ if TYPE_CHECKING:
 
 
 @final
-class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # type:ignore[name-defined]
+class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     """Handler for PyTorch JIT based model.
 
     Currently torch.jit.ScriptModule based classes are supported.
@@ -41,25 +41,25 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
     def can_handle(
         cls,
         model: model_types.SupportedModelType,
-    ) -> TypeGuard["torch.jit.ScriptModule"]:  # type:ignore[name-defined]
+    ) -> TypeGuard["torch.jit.ScriptModule"]:
         return type_utils.LazyType("torch.jit.ScriptModule").isinstance(model)
 
     @classmethod
     def cast_model(
         cls,
         model: model_types.SupportedModelType,
-    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+    ) -> "torch.jit.ScriptModule":
         import torch
 
-        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(model, torch.jit.ScriptModule)
 
-        return cast(torch.jit.ScriptModule, model)  # type:ignore[name-defined]
+        return cast(torch.jit.ScriptModule, model)
 
     @classmethod
     def save_model(
         cls,
         name: str,
-        model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        model: "torch.jit.ScriptModule",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
@@ -72,7 +72,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
 
         import torch
 
-        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(model, torch.jit.ScriptModule)
 
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -111,7 +111,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
         with open(os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
-            torch.jit.save(model, f)  # type:ignore[no-untyped-call, attr-defined]
+            torch.jit.save(model, f)  # type:ignore[no-untyped-call]
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
@@ -133,7 +133,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         **kwargs: Unpack[model_types.TorchScriptLoadOptions],
-    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+    ) -> "torch.jit.ScriptModule":
         import torch
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
@@ -141,10 +141,10 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_blob_metadata = model_blobs_metadata[name]
         model_blob_filename = model_blob_metadata.path
         with open(os.path.join(model_blob_path, model_blob_filename), "rb") as f:
-            m = torch.jit.load(  # type:ignore[no-untyped-call, attr-defined]
+            m = torch.jit.load(  # type:ignore[no-untyped-call]
                 f, map_location="cuda" if kwargs.get("use_gpu", False) else "cpu"
             )
-        assert isinstance(m, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(m, torch.jit.ScriptModule)
 
         if kwargs.get("use_gpu", False):
             m = m.cuda()
@@ -154,7 +154,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
     @classmethod
     def convert_as_custom_model(
         cls,
-        raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        raw_model: "torch.jit.ScriptModule",
         model_meta: model_meta_api.ModelMetadata,
         background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.TorchScriptLoadOptions],
@@ -162,11 +162,11 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         from snowflake.ml.model import custom_model
 
         def _create_custom_model(
-            raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+            raw_model: "torch.jit.ScriptModule",
             model_meta: model_meta_api.ModelMetadata,
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
-                raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+                raw_model: "torch.jit.ScriptModule",
                 signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:

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

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

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

@@ -58,11 +58,16 @@ class XgboostModelBlobOptions(BaseModelBlobOptions):
     xgb_estimator_type: Required[str]
 
 
+class TensorflowModelBlobOptions(BaseModelBlobOptions):
+    is_keras_model: Required[bool]
+
+
 ModelBlobOptions = Union[
     BaseModelBlobOptions,
     HuggingFacePipelineModelBlobOptions,
     MLFlowModelBlobOptions,
     XgboostModelBlobOptions,
+    TensorflowModelBlobOptions,
 ]
 
 

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

@@ -1,10 +1,2 @@
-REQUIREMENTS = [
-    "absl-py>=0.15,<2",
-    "anyio>=3.5.0,<4",
-    "numpy>=1.23,<2",
-    "packaging>=20.9,<24",
-    "pandas>=1.0.0,<3",
-    "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.17.0,<2",
-    "typing-extensions>=4.1.0,<5"
-]
+REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2022.11,<2024', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2022.11,<2024', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', '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,<4', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2022.11,<2024', 'importlib_resources>=6.1.1, <7', 'lightgbm>=4.1.0, <5', 'mlflow>=2.1.0,<2.4', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pytimeparse>=1.1.8,<2', 'pytorch>=2.0.1,<2.3.0', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2022.11,<2024', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'sentence-transformers>=2.2.2,<3', 'sentencepiece>=0.1.95,<1', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'tensorflow>=2.10,<3', 'tokenizers>=0.10,<1', 'torchdata>=0.4,<1', 'transformers>=4.32.1,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']

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

@@ -17,6 +17,8 @@ _SNOWML_INFERENCE_ALTERNATIVE_DEPENDENCIES = [
     for r in _snowml_inference_alternative_requirements.REQUIREMENTS
 ]
 
+PACKAGES_NOT_ALLOWED_IN_WAREHOUSE = ["snowflake-connector-python", "pyarrow"]
+
 
 class ModelRuntime:
     """Class to represent runtime in a model, which controls the runtime and version, imports and dependencies.
@@ -61,15 +63,8 @@ class ModelRuntime:
             ],
         )
 
-        if not is_warehouse and self.embed_local_ml_library:
-            self.runtime_env.include_if_absent(
-                [
-                    model_env.ModelDependency(
-                        requirement="pyarrow",
-                        pip_name="pyarrow",
-                    )
-                ],
-            )
+        if is_warehouse and self.embed_local_ml_library:
+            self.runtime_env.remove_if_present_conda(PACKAGES_NOT_ALLOWED_IN_WAREHOUSE)
 
         if is_gpu:
             self.runtime_env.generate_env_for_cuda()

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

@@ -84,7 +84,7 @@ def get_model_task_lightgbm(model: Union["lightgbm.Booster", "lightgbm.LGBMModel
     if type_utils.LazyType("lightgbm.Booster").isinstance(model):
         model_task = model.params["objective"]  # type: ignore[attr-defined]
     elif hasattr(model, "objective_"):
-        model_task = model.objective_
+        model_task = model.objective_  # type: ignore[assignment]
     if model_task in _BINARY_CLASSIFICATION_OBJECTIVES:
         return type_hints.Task.TABULAR_BINARY_CLASSIFICATION
     if model_task in _MULTI_CLASSIFICATION_OBJECTIVES:

snowflake/ml/model/_signatures/core.py

@@ -14,10 +14,12 @@ from typing import (
     Type,
     Union,
     final,
+    get_args,
 )
 
 import numpy as np
 import numpy.typing as npt
+import pandas as pd
 
 import snowflake.snowpark.types as spt
 from snowflake.ml._internal.exceptions import (
@@ -29,6 +31,21 @@ if TYPE_CHECKING:
     import mlflow
     import torch
 
+PandasExtensionTypes = Union[
+    pd.Int8Dtype,
+    pd.Int16Dtype,
+    pd.Int32Dtype,
+    pd.Int64Dtype,
+    pd.UInt8Dtype,
+    pd.UInt16Dtype,
+    pd.UInt32Dtype,
+    pd.UInt64Dtype,
+    pd.Float32Dtype,
+    pd.Float64Dtype,
+    pd.BooleanDtype,
+    pd.StringDtype,
+]
+
 
 class DataType(Enum):
     def __init__(self, value: str, snowpark_type: Type[spt.DataType], numpy_type: npt.DTypeLike) -> None:
@@ -67,11 +84,11 @@ class DataType(Enum):
         return f"DataType.{self.name}"
 
     @classmethod
-    def from_numpy_type(cls, np_type: npt.DTypeLike) -> "DataType":
+    def from_numpy_type(cls, input_type: Union[npt.DTypeLike, PandasExtensionTypes]) -> "DataType":
         """Translate numpy dtype to DataType for signature definition.
 
         Args:
-            np_type: The numpy dtype.
+            input_type: The numpy dtype or Pandas Extension Dtype
 
         Raises:
             SnowflakeMLException: NotImplementedError: Raised when the given numpy type is not supported.
@@ -79,6 +96,10 @@ class DataType(Enum):
         Returns:
             Corresponding DataType.
         """
+        # To support pandas extension dtype
+        if isinstance(input_type, get_args(PandasExtensionTypes)):
+            input_type = input_type.type
+
         np_to_snowml_type_mapping = {i._numpy_type: i for i in DataType}
 
         # Add datetime types:
@@ -88,12 +109,12 @@ class DataType(Enum):
             np_to_snowml_type_mapping[f"datetime64[{res}]"] = DataType.TIMESTAMP_NTZ
 
         for potential_type in np_to_snowml_type_mapping.keys():
-            if np.can_cast(np_type, potential_type, casting="no"):
+            if np.can_cast(input_type, potential_type, casting="no"):
                 # This is used since the same dtype might represented in different ways.
                 return np_to_snowml_type_mapping[potential_type]
         raise snowml_exceptions.SnowflakeMLException(
             error_code=error_codes.NOT_IMPLEMENTED,
-            original_exception=NotImplementedError(f"Type {np_type} is not supported as a DataType."),
+            original_exception=NotImplementedError(f"Type {input_type} is not supported as a DataType."),
         )
 
     @classmethod
@@ -212,6 +233,7 @@ class FeatureSpec(BaseFeatureSpec):
         name: str,
         dtype: DataType,
         shape: Optional[Tuple[int, ...]] = None,
+        nullable: bool = True,
     ) -> None:
         """
         Initialize a feature.
@@ -219,6 +241,7 @@ class FeatureSpec(BaseFeatureSpec):
         Args:
             name: Name of the feature.
             dtype: Type of the elements in the feature.
+            nullable: Whether the feature is nullable. Defaults to True.
             shape: Used to represent scalar feature, 1-d feature list,
                 or n-d tensor. Use -1 to represent variable length. Defaults to None.
 
@@ -227,6 +250,7 @@ class FeatureSpec(BaseFeatureSpec):
                     - (2,): 1d list with a fixed length of 2.
                     - (-1,): 1d list with variable length, used for ragged tensor representation.
                     - (d1, d2, d3): 3d tensor.
+            nullable: Whether the feature is nullable. Defaults to True.
 
         Raises:
             SnowflakeMLException: TypeError: When the dtype input type is incorrect.
@@ -248,6 +272,8 @@ class FeatureSpec(BaseFeatureSpec):
             )
         self._shape = shape
 
+        self._nullable = nullable
+
     def as_snowpark_type(self) -> spt.DataType:
         result_type = self._dtype.as_snowpark_type()
         if not self._shape:
@@ -256,13 +282,34 @@ class FeatureSpec(BaseFeatureSpec):
             result_type = spt.ArrayType(result_type)
         return result_type
 
-    def as_dtype(self) -> Union[npt.DTypeLike, str]:
+    def as_dtype(self) -> Union[npt.DTypeLike, str, PandasExtensionTypes]:
         """Convert to corresponding local Type."""
+
         if not self._shape:
             # scalar dtype: use keys from `np.sctypeDict` to prevent unit-less dtype 'datetime64'
             if "datetime64" in self._dtype._value:
                 return self._dtype._value
-            return self._dtype._numpy_type
+
+            np_type = self._dtype._numpy_type
+            if self._nullable:
+                np_to_pd_dtype_mapping = {
+                    np.int8: pd.Int8Dtype(),
+                    np.int16: pd.Int16Dtype(),
+                    np.int32: pd.Int32Dtype(),
+                    np.int64: pd.Int64Dtype(),
+                    np.uint8: pd.UInt8Dtype(),
+                    np.uint16: pd.UInt16Dtype(),
+                    np.uint32: pd.UInt32Dtype(),
+                    np.uint64: pd.UInt64Dtype(),
+                    np.float32: pd.Float32Dtype(),
+                    np.float64: pd.Float64Dtype(),
+                    np.bool_: pd.BooleanDtype(),
+                    np.str_: pd.StringDtype(),
+                }
+
+                return np_to_pd_dtype_mapping.get(np_type, np_type)  # type: ignore[arg-type]
+
+            return np_type
         return np.object_
 
     def __eq__(self, other: object) -> bool:
@@ -273,7 +320,10 @@ class FeatureSpec(BaseFeatureSpec):
 
     def __repr__(self) -> str:
         shape_str = f", shape={repr(self._shape)}" if self._shape else ""
-        return f"FeatureSpec(dtype={repr(self._dtype)}, name={repr(self._name)}{shape_str})"
+        return (
+            f"FeatureSpec(dtype={repr(self._dtype)}, "
+            f"name={repr(self._name)}{shape_str}, nullable={repr(self._nullable)})"
+        )
 
     def to_dict(self) -> Dict[str, Any]:
         """Serialize the feature group into a dict.
@@ -281,10 +331,7 @@ class FeatureSpec(BaseFeatureSpec):
         Returns:
             A dict that serializes the feature group.
         """
-        base_dict: Dict[str, Any] = {
-            "type": self._dtype.name,
-            "name": self._name,
-        }
+        base_dict: Dict[str, Any] = {"type": self._dtype.name, "name": self._name, "nullable": self._nullable}
         if self._shape is not None:
             base_dict["shape"] = self._shape
         return base_dict
@@ -304,7 +351,9 @@ class FeatureSpec(BaseFeatureSpec):
         if shape:
             shape = tuple(shape)
         type = DataType[input_dict["type"]]
-        return FeatureSpec(name=name, dtype=type, shape=shape)
+        # If nullable is not provided, default to False for backward compatibility.
+        nullable = input_dict.get("nullable", False)
+        return FeatureSpec(name=name, dtype=type, shape=shape, nullable=nullable)
 
     @classmethod
     def from_mlflow_spec(
@@ -475,10 +524,8 @@ class ModelSignature:
         sig_outs = loaded["outputs"]
         sig_inputs = loaded["inputs"]
 
-        deserialize_spec: Callable[[Dict[str, Any]], BaseFeatureSpec] = (
-            lambda sig_spec: FeatureGroupSpec.from_dict(sig_spec)
-            if "feature_group" in sig_spec
-            else FeatureSpec.from_dict(sig_spec)
+        deserialize_spec: Callable[[Dict[str, Any]], BaseFeatureSpec] = lambda sig_spec: (
+            FeatureGroupSpec.from_dict(sig_spec) if "feature_group" in sig_spec else FeatureSpec.from_dict(sig_spec)
         )
 
         return ModelSignature(