snowflake-ml-python 1.8.2__py3-none-any.whl → 1.8.4__py3-none-any.whl

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

@@ -6,31 +6,26 @@ import zipfile
 from contextlib import contextmanager
 from datetime import datetime
 from types import ModuleType
-from typing import Any, Dict, Generator, List, Optional, TypedDict
+from typing import Any, Generator, Optional, TypedDict
 
 import cloudpickle
 import yaml
 from packaging import requirements, version
 from typing_extensions import Required
 
-from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml import version as snowml_version
+from snowflake.ml._internal import env_utils, file_utils
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_meta import (
-    _packaging_requirements,
-    model_blob_meta,
-    model_meta_schema,
-)
+from snowflake.ml.model._packager.model_meta import model_blob_meta, model_meta_schema
 from snowflake.ml.model._packager.model_meta_migrator import migrator_plans
 from snowflake.ml.model._packager.model_runtime import model_runtime
 
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
 
-_PACKAGING_REQUIREMENTS = [
-    str(env_utils.get_package_spec_with_supported_ops_only(requirements.Requirement(r)))
-    for r in _packaging_requirements.REQUIREMENTS
-]
+_PACKAGING_REQUIREMENTS = ["cloudpickle"]
+
 _SNOWFLAKE_PKG_NAME = "snowflake"
 _SNOWFLAKE_ML_PKG_NAME = f"{_SNOWFLAKE_PKG_NAME}.ml"
 
@@ -41,15 +36,16 @@ def create_model_metadata(
     model_dir_path: str,
     name: str,
     model_type: model_types.SupportedModelHandlerType,
-    signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-    function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
-    metadata: Optional[Dict[str, str]] = None,
-    code_paths: Optional[List[str]] = None,
-    ext_modules: Optional[List[ModuleType]] = None,
-    conda_dependencies: Optional[List[str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    artifact_repository_map: Optional[Dict[str, str]] = None,
-    target_platforms: Optional[List[model_types.TargetPlatform]] = None,
+    signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
+    function_properties: Optional[dict[str, dict[str, Any]]] = None,
+    metadata: Optional[dict[str, str]] = None,
+    code_paths: Optional[list[str]] = None,
+    ext_modules: Optional[list[ModuleType]] = None,
+    conda_dependencies: Optional[list[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    artifact_repository_map: Optional[dict[str, str]] = None,
+    resource_constraint: Optional[dict[str, str]] = None,
+    target_platforms: Optional[list[model_types.TargetPlatform]] = None,
     python_version: Optional[str] = None,
     task: model_types.Task = model_types.Task.UNKNOWN,
     **kwargs: Any,
@@ -69,7 +65,9 @@ def create_model_metadata(
         ext_modules: List of names of modules that need to be pickled with the model. Defaults to None.
         conda_dependencies: List of conda requirements for running the model. Defaults to None.
         pip_requirements: List of pip Python packages requirements for running the model. Defaults to None.
-        artifact_repository_map: A dict mapping from package channel to artifact repository name.
+        artifact_repository_map: A dict mapping from package channel to artifact repository name (e.g.
+            {'pip': 'snowflake.snowpark.pypi_shared_repository'}).
+        resource_constraint: Mapping of resource constraint keys and values, e.g. {"architecture": "x86"}.
         target_platforms: List of target platforms to run the model.
         python_version: A string of python version where model is run. Used for user override. If specified as None,
             current version would be captured. Defaults to None.
@@ -108,13 +106,14 @@ def create_model_metadata(
         conda_dependencies=conda_dependencies,
         pip_requirements=pip_requirements,
         artifact_repository_map=artifact_repository_map,
+        resource_constraint=resource_constraint,
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
         prefer_pip=prefer_pip,
     )
 
     if embed_local_ml_library:
-        env.snowpark_ml_version = f"{snowml_env.VERSION}+{file_utils.hash_directory(path_to_copy)}"
+        env.snowpark_ml_version = f"{snowml_version.VERSION}+{file_utils.hash_directory(path_to_copy)}"
 
     model_meta = ModelMetadata(
         name=name,
@@ -156,9 +155,10 @@ def create_model_metadata(
 
 def _create_env_for_model_metadata(
     *,
-    conda_dependencies: Optional[List[str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    artifact_repository_map: Optional[Dict[str, str]] = None,
+    conda_dependencies: Optional[list[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    artifact_repository_map: Optional[dict[str, str]] = None,
+    resource_constraint: Optional[dict[str, str]] = None,
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
     prefer_pip: bool = False,
@@ -169,8 +169,9 @@ def _create_env_for_model_metadata(
     env.conda_dependencies = conda_dependencies  # type: ignore[assignment]
     env.pip_requirements = pip_requirements  # type: ignore[assignment]
     env.artifact_repository_map = artifact_repository_map
+    env.resource_constraint = resource_constraint
     env.python_version = python_version  # type: ignore[assignment]
-    env.snowpark_ml_version = snowml_env.VERSION
+    env.snowpark_ml_version = snowml_version.VERSION
 
     requirements_to_add = _PACKAGING_REQUIREMENTS
 
@@ -242,20 +243,20 @@ class ModelMetadata:
         name: str,
         env: model_env.ModelEnv,
         model_type: model_types.SupportedModelHandlerType,
-        runtimes: Optional[Dict[str, model_runtime.ModelRuntime]] = None,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
-        function_properties: Optional[Dict[str, Dict[str, Any]]] = None,
-        user_files: Optional[Dict[str, List[str]]] = None,
-        metadata: Optional[Dict[str, str]] = None,
+        runtimes: Optional[dict[str, model_runtime.ModelRuntime]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
+        function_properties: Optional[dict[str, dict[str, Any]]] = None,
+        user_files: Optional[dict[str, list[str]]] = None,
+        metadata: Optional[dict[str, str]] = None,
         creation_timestamp: Optional[str] = None,
         min_snowpark_ml_version: Optional[str] = None,
-        models: Optional[Dict[str, model_blob_meta.ModelBlobMeta]] = None,
+        models: Optional[dict[str, model_blob_meta.ModelBlobMeta]] = None,
         original_metadata_version: Optional[str] = model_meta_schema.MODEL_METADATA_VERSION,
         task: model_types.Task = model_types.Task.UNKNOWN,
         explain_algorithm: Optional[model_meta_schema.ModelExplainAlgorithm] = None,
     ) -> None:
         self.name = name
-        self.signatures: Dict[str, model_signature.ModelSignature] = dict()
+        self.signatures: dict[str, model_signature.ModelSignature] = dict()
         if signatures:
             self.signatures = signatures
         self.function_properties = function_properties or {}
@@ -270,7 +271,7 @@ class ModelMetadata:
             else model_meta_schema.MODEL_METADATA_MIN_SNOWPARK_ML_VERSION
         )
 
-        self.models: Dict[str, model_blob_meta.ModelBlobMeta] = dict()
+        self.models: dict[str, model_blob_meta.ModelBlobMeta] = dict()
         if models:
             self.models = models
 
@@ -291,7 +292,7 @@ class ModelMetadata:
         self._min_snowpark_ml_version = max(self._min_snowpark_ml_version, parsed_min_snowpark_ml_version)
 
     @property
-    def runtimes(self) -> Dict[str, model_runtime.ModelRuntime]:
+    def runtimes(self) -> dict[str, model_runtime.ModelRuntime]:
         if self._runtimes and "cpu" in self._runtimes:
             return self._runtimes
         runtimes = {
@@ -358,11 +359,11 @@ class ModelMetadata:
 
         loaded_meta_min_snowpark_ml_version = loaded_meta.get("min_snowpark_ml_version", None)
         if not loaded_meta_min_snowpark_ml_version or (
-            version.parse(loaded_meta_min_snowpark_ml_version) > version.parse(snowml_env.VERSION)
+            version.parse(loaded_meta_min_snowpark_ml_version) > version.parse(snowml_version.VERSION)
         ):
             raise RuntimeError(
                 f"The minimal version required to load the model is {loaded_meta_min_snowpark_ml_version}, "
-                f"while current version of Snowpark ML library is {snowml_env.VERSION}."
+                f"while current version of Snowpark ML library is {snowml_version.VERSION}."
             )
         return model_meta_schema.ModelMetadataDict(
             creation_timestamp=loaded_meta["creation_timestamp"],
@@ -405,7 +406,7 @@ class ModelMetadata:
         env = model_env.ModelEnv()
         env.load_from_dict(pathlib.Path(model_dir_path), model_dict["env"])
 
-        runtimes: Optional[Dict[str, model_runtime.ModelRuntime]]
+        runtimes: Optional[dict[str, model_runtime.ModelRuntime]]
         if model_dict.get("runtimes", None):
             runtimes = {
                 name: model_runtime.ModelRuntime.load(pathlib.Path(model_dir_path), name, env, runtime_dict)

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

@@ -1,7 +1,7 @@
 # This files contains schema definition of what will be written into model.yml
 # Changing this file should lead to a change of the schema version.
 from enum import Enum
-from typing import Any, Dict, List, Optional, TypedDict, Union
+from typing import Any, Optional, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
 
@@ -18,18 +18,20 @@ class FunctionProperties(Enum):
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
-    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
+    artifact_repository_map: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelRuntimeDict(TypedDict):
-    imports: Required[List[str]]
+    imports: Required[list[str]]
     dependencies: Required[ModelRuntimeDependenciesDict]
+    resource_constraint: NotRequired[Optional[dict[str, str]]]
 
 
 class ModelEnvDict(TypedDict):
     conda: Required[str]
     pip: Required[str]
-    artifact_repository_map: NotRequired[Optional[Dict[str, str]]]
+    artifact_repository_map: NotRequired[Optional[dict[str, str]]]
+    resource_constraint: NotRequired[Optional[dict[str, str]]]
     python_version: Required[str]
     cuda_version: NotRequired[Optional[str]]
     snowpark_ml_version: Required[str]
@@ -61,6 +63,7 @@ class MLFlowModelBlobOptions(BaseModelBlobOptions):
 
 class XgboostModelBlobOptions(BaseModelBlobOptions):
     xgb_estimator_type: Required[str]
+    enable_categorical: NotRequired[bool]
 
 
 class PyTorchModelBlobOptions(BaseModelBlobOptions):
@@ -102,25 +105,25 @@ class ModelBlobMetadataDict(TypedDict):
     model_type: Required[type_hints.SupportedModelHandlerType]
     path: Required[str]
     handler_version: Required[str]
-    function_properties: NotRequired[Dict[str, Dict[str, Any]]]
-    artifacts: NotRequired[Dict[str, str]]
+    function_properties: NotRequired[dict[str, dict[str, Any]]]
+    artifacts: NotRequired[dict[str, str]]
     options: NotRequired[ModelBlobOptions]
 
 
 class ModelMetadataDict(TypedDict):
     creation_timestamp: Required[str]
     env: Required[ModelEnvDict]
-    runtimes: NotRequired[Dict[str, ModelRuntimeDict]]
-    metadata: NotRequired[Optional[Dict[str, str]]]
+    runtimes: NotRequired[dict[str, ModelRuntimeDict]]
+    metadata: NotRequired[Optional[dict[str, str]]]
     model_type: Required[type_hints.SupportedModelHandlerType]
-    models: Required[Dict[str, ModelBlobMetadataDict]]
+    models: Required[dict[str, ModelBlobMetadataDict]]
     name: Required[str]
-    signatures: Required[Dict[str, Dict[str, Any]]]
+    signatures: Required[dict[str, dict[str, Any]]]
     version: Required[str]
     min_snowpark_ml_version: Required[str]
     task: Required[str]
     explainability: NotRequired[Optional[ExplainabilityMetadataDict]]
-    function_properties: NotRequired[Dict[str, Dict[str, Any]]]
+    function_properties: NotRequired[dict[str, dict[str, Any]]]
 
 
 class ModelExplainAlgorithm(Enum):

snowflake/ml/model/_packager/model_meta_migrator/base_migrator.py

@@ -1,6 +1,6 @@
 import copy
 from abc import abstractmethod
-from typing import Any, Dict, Protocol, final
+from typing import Any, Protocol, final
 
 from snowflake.ml._internal import migrator_utils
 
@@ -11,13 +11,13 @@ class _BaseModelMetaMigratorProtocol(Protocol):
 
     @staticmethod
     @abstractmethod
-    def upgrade(original_meta_dict: Dict[str, Any]) -> Dict[str, Any]:
+    def upgrade(original_meta_dict: dict[str, Any]) -> dict[str, Any]:
         raise NotImplementedError
 
 
 class BaseModelMetaMigrator(_BaseModelMetaMigratorProtocol):
     @final
-    def try_upgrade(self, original_meta_dict: Dict[str, Any]) -> Dict[str, Any]:
+    def try_upgrade(self, original_meta_dict: dict[str, Any]) -> dict[str, Any]:
         loaded_meta_version = original_meta_dict.get("version", None)
         if not loaded_meta_version or str(loaded_meta_version) != self.source_version:
             raise NotImplementedError(

snowflake/ml/model/_packager/model_meta_migrator/migrator_plans.py

@@ -1,14 +1,14 @@
-from typing import Any, Dict, Type
+from typing import Any
 
 from snowflake.ml.model._packager.model_meta import model_meta_schema
 from snowflake.ml.model._packager.model_meta_migrator import base_migrator, migrator_v1
 
-MODEL_META_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelMetaMigrator]] = {
+MODEL_META_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelMetaMigrator]] = {
     "1": migrator_v1.MetaMigrator_v1,
 }
 
 
-def migrate_metadata(loaded_meta: Dict[str, Any]) -> Dict[str, Any]:
+def migrate_metadata(loaded_meta: dict[str, Any]) -> dict[str, Any]:
     loaded_meta_version = str(loaded_meta.get("version", None))
     while loaded_meta_version != model_meta_schema.MODEL_METADATA_VERSION:
         if loaded_meta_version not in MODEL_META_MIGRATOR_PLANS.keys():

snowflake/ml/model/_packager/model_meta_migrator/migrator_v1.py

@@ -1,8 +1,8 @@
-from typing import Any, Dict
+from typing import Any
 
 from packaging import requirements, version
 
-from snowflake.ml._internal import env as snowml_env
+from snowflake.ml import version as snowml_version
 from snowflake.ml.model._packager.model_meta_migrator import base_migrator
 
 
@@ -11,7 +11,7 @@ class MetaMigrator_v1(base_migrator.BaseModelMetaMigrator):
     target_version = "2023-12-01"
 
     @staticmethod
-    def upgrade(original_meta_dict: Dict[str, Any]) -> Dict[str, Any]:
+    def upgrade(original_meta_dict: dict[str, Any]) -> dict[str, Any]:
         loaded_python_version = version.parse(original_meta_dict["python_version"])
         if original_meta_dict.get("local_ml_library_version", None):
             loaded_lib_version = str(version.parse(original_meta_dict["local_ml_library_version"]))
@@ -24,7 +24,7 @@ class MetaMigrator_v1(base_migrator.BaseModelMetaMigrator):
                 None,
             )
             if lib_spec_str is None:
-                loaded_lib_version = snowml_env.VERSION
+                loaded_lib_version = snowml_version.VERSION
             loaded_lib_version = list(requirements.Requirement(str(lib_spec_str)).specifier)[0].version
 
         return dict(

snowflake/ml/model/_packager/model_packager.py

@@ -1,6 +1,6 @@
 import os
 from types import ModuleType
-from typing import Dict, List, Optional
+from typing import Optional
 
 from absl import logging
 
@@ -38,16 +38,17 @@ class ModelPackager:
         *,
         name: str,
         model: model_types.SupportedModelType,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        metadata: Optional[Dict[str, str]] = None,
-        conda_dependencies: Optional[List[str]] = None,
-        pip_requirements: Optional[List[str]] = None,
-        artifact_repository_map: Optional[Dict[str, str]] = None,
-        target_platforms: Optional[List[model_types.TargetPlatform]] = None,
+        metadata: Optional[dict[str, str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
+        artifact_repository_map: Optional[dict[str, str]] = None,
+        resource_constraint: Optional[dict[str, str]] = None,
+        target_platforms: Optional[list[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
-        code_paths: Optional[List[str]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
+        code_paths: Optional[list[str]] = None,
         options: model_types.ModelSaveOption,
         task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:
@@ -76,6 +77,7 @@ class ModelPackager:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             artifact_repository_map=artifact_repository_map,
+            resource_constraint=resource_constraint,
             python_version=python_version,
             task=task,
             target_platforms=target_platforms,

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

@@ -1 +1,32 @@
-REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<5', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0,<3', 'cryptography', 'fsspec>=2024.6.1,<2026', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2024.6.1,<2026', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.12.0,<4', 'snowflake-snowpark-python>=1.17.0,<2,!=1.26.0', 'snowflake.core>=1.0.2,<2', 'sqlparse>=0.4,<1', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
+# DO NOT EDIT!
+# Generate by running 'bazel run --config=pre_build //bazel/requirements:sync_requirements'
+
+REQUIREMENTS = [
+    "absl-py>=0.15,<2",
+    "aiohttp!=4.0.0a0, !=4.0.0a1",
+    "anyio>=3.5.0,<5",
+    "cachetools>=3.1.1,<6",
+    "cloudpickle>=2.0.0",
+    "cryptography",
+    "fsspec>=2024.6.1,<2026",
+    "importlib_resources>=6.1.1, <7",
+    "numpy>=1.23,<2",
+    "packaging>=20.9,<25",
+    "pandas>=1.0.0,<3",
+    "pyarrow",
+    "pydantic>=2.8.2, <3",
+    "pyjwt>=2.0.0, <3",
+    "pytimeparse>=1.1.8,<2",
+    "pyyaml>=6.0,<7",
+    "requests",
+    "retrying>=1.3.3,<2",
+    "s3fs>=2024.6.1,<2026",
+    "scikit-learn<1.6",
+    "scipy>=1.9,<2",
+    "shap>=0.46.0,<1",
+    "snowflake-connector-python>=3.14.0,<4",
+    "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0",
+    "snowflake.core>=1.0.2,<2",
+    "sqlparse>=0.4,<1",
+    "typing-extensions>=4.1.0,<5",
+]

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

@@ -1,7 +1,7 @@
 import copy
 import pathlib
 import warnings
-from typing import List, Optional
+from typing import Optional
 
 from packaging import requirements
 
@@ -37,7 +37,7 @@ class ModelRuntime:
         self,
         name: str,
         env: model_env.ModelEnv,
-        imports: Optional[List[str]] = None,
+        imports: Optional[list[str]] = None,
         is_warehouse: bool = False,
         is_gpu: bool = False,
         loading_from_file: bool = False,
@@ -102,6 +102,7 @@ class ModelRuntime:
                 if env_dict.get("artifact_repository_map") is not None
                 else {},
             ),
+            resource_constraint=env_dict["resource_constraint"],
         )
 
     @staticmethod
@@ -116,6 +117,7 @@ class ModelRuntime:
         env.cuda_version = meta_env.cuda_version
         env.snowpark_ml_version = meta_env.snowpark_ml_version
         env.artifact_repository_map = meta_env.artifact_repository_map
+        env.resource_constraint = meta_env.resource_constraint
 
         conda_env_rel_path = pathlib.PurePosixPath(loaded_dict["dependencies"]["conda"])
         pip_requirements_rel_path = pathlib.PurePosixPath(loaded_dict["dependencies"]["pip"])

snowflake/ml/model/_signatures/core.py

@@ -6,12 +6,8 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
-    Dict,
-    List,
     Optional,
     Sequence,
-    Tuple,
-    Type,
     Union,
     final,
     get_args,
@@ -48,7 +44,7 @@ PandasExtensionTypes = Union[
 
 
 class DataType(Enum):
-    def __init__(self, value: str, snowpark_type: Type[spt.DataType], numpy_type: npt.DTypeLike) -> None:
+    def __init__(self, value: str, snowpark_type: type[spt.DataType], numpy_type: npt.DTypeLike) -> None:
         self._value = value
         self._snowpark_type = snowpark_type
         self._numpy_type = numpy_type
@@ -159,7 +155,7 @@ class DataType(Enum):
         else:
             actual_sp_type = snowpark_type
 
-        snowpark_to_snowml_type_mapping: Dict[Type[spt.DataType], DataType] = {
+        snowpark_to_snowml_type_mapping: dict[type[spt.DataType], DataType] = {
             i._snowpark_type: i
             for i in DataType
             # We by default infer as signed integer.
@@ -199,7 +195,7 @@ class DataType(Enum):
 class BaseFeatureSpec(ABC):
     """Abstract Class for specification of a feature."""
 
-    def __init__(self, name: str, shape: Optional[Tuple[int, ...]]) -> None:
+    def __init__(self, name: str, shape: Optional[tuple[int, ...]]) -> None:
         self._name = name
 
         if shape and not isinstance(shape, tuple):
@@ -218,23 +214,19 @@ class BaseFeatureSpec(ABC):
     @abstractmethod
     def as_snowpark_type(self) -> spt.DataType:
         """Convert to corresponding Snowpark Type."""
-        pass
 
     @abstractmethod
     def as_dtype(self, force_numpy_dtype: bool = False) -> Union[npt.DTypeLike, str, PandasExtensionTypes]:
         """Convert to corresponding local Type."""
-        pass
 
     @abstractmethod
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Serialization"""
-        pass
 
     @classmethod
     @abstractmethod
-    def from_dict(self, input_dict: Dict[str, Any]) -> "BaseFeatureSpec":
+    def from_dict(self, input_dict: dict[str, Any]) -> "BaseFeatureSpec":
         """Deserialization"""
-        pass
 
 
 class FeatureSpec(BaseFeatureSpec):
@@ -244,7 +236,7 @@ class FeatureSpec(BaseFeatureSpec):
         self,
         name: str,
         dtype: DataType,
-        shape: Optional[Tuple[int, ...]] = None,
+        shape: Optional[tuple[int, ...]] = None,
         nullable: bool = True,
     ) -> None:
         """
@@ -330,19 +322,19 @@ class FeatureSpec(BaseFeatureSpec):
             f"name={repr(self._name)}{shape_str}, nullable={repr(self._nullable)})"
         )
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Serialize the feature group into a dict.
 
         Returns:
             A dict that serializes the feature group.
         """
-        base_dict: Dict[str, Any] = {"type": self._dtype.name, "name": self._name, "nullable": self._nullable}
+        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
 
     @classmethod
-    def from_dict(cls, input_dict: Dict[str, Any]) -> "FeatureSpec":
+    def from_dict(cls, input_dict: dict[str, Any]) -> "FeatureSpec":
         """Deserialize the feature specification from a dict.
 
         Args:
@@ -391,7 +383,7 @@ class FeatureSpec(BaseFeatureSpec):
 class FeatureGroupSpec(BaseFeatureSpec):
     """Specification of a group of features in Snowflake native model packaging."""
 
-    def __init__(self, name: str, specs: List[BaseFeatureSpec], shape: Optional[Tuple[int, ...]] = None) -> None:
+    def __init__(self, name: str, specs: list[BaseFeatureSpec], shape: Optional[tuple[int, ...]] = None) -> None:
         """Initialize a feature group.
 
         Args:
@@ -458,19 +450,19 @@ class FeatureGroupSpec(BaseFeatureSpec):
     def as_dtype(self, force_numpy_dtype: bool = False) -> Union[npt.DTypeLike, str, PandasExtensionTypes]:
         return np.object_
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Serialize the feature group into a dict.
 
         Returns:
             A dict that serializes the feature group.
         """
-        base_dict: Dict[str, Any] = {"name": self._name, "specs": [s.to_dict() for s in self._specs]}
+        base_dict: dict[str, Any] = {"name": self._name, "specs": [s.to_dict() for s in self._specs]}
         if self._shape is not None:
             base_dict["shape"] = self._shape
         return base_dict
 
     @classmethod
-    def from_dict(cls, input_dict: Dict[str, Any]) -> "FeatureGroupSpec":
+    def from_dict(cls, input_dict: dict[str, Any]) -> "FeatureGroupSpec":
         """Deserialize the feature group from a dict.
 
         Args:
@@ -520,7 +512,7 @@ class ModelSignature:
         else:
             return False
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """Generate a dict to represent the whole signature.
 
         Returns:
@@ -533,7 +525,7 @@ class ModelSignature:
         }
 
     @classmethod
-    def from_dict(cls, loaded: Dict[str, Any]) -> "ModelSignature":
+    def from_dict(cls, loaded: dict[str, Any]) -> "ModelSignature":
         """Create a signature given the dict containing specifications of children features and feature groups.
 
         Args:
@@ -545,7 +537,7 @@ class ModelSignature:
         sig_outs = loaded["outputs"]
         sig_inputs = loaded["inputs"]
 
-        deserialize_spec: Callable[[Dict[str, Any]], BaseFeatureSpec] = lambda sig_spec: (
+        deserialize_spec: Callable[[dict[str, Any]], BaseFeatureSpec] = lambda sig_spec: (
             FeatureGroupSpec.from_dict(sig_spec) if "specs" in sig_spec else FeatureSpec.from_dict(sig_spec)
         )
 

snowflake/ml/model/_signatures/dmatrix_handler.py

@@ -1,4 +1,4 @@
-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
@@ -45,7 +45,7 @@ class XGBoostDMatrixHandler(base_handler.BaseDataHandler["xgboost.DMatrix"]):
     @staticmethod
     def infer_signature(data: "xgboost.DMatrix", role: Literal["input", "output"]) -> Sequence[core.BaseFeatureSpec]:
         feature_prefix = f"{XGBoostDMatrixHandler.FEATURE_PREFIX}_"
-        features: List[core.BaseFeatureSpec] = []
+        features: list[core.BaseFeatureSpec] = []
         role_prefix = (
             XGBoostDMatrixHandler.INPUT_PREFIX if role == "input" else XGBoostDMatrixHandler.OUTPUT_PREFIX
         ) + "_"
@@ -81,8 +81,16 @@ class XGBoostDMatrixHandler(base_handler.BaseDataHandler["xgboost.DMatrix"]):
     ) -> "xgboost.DMatrix":
         import xgboost as xgb
 
+        enable_categorical = False
+        for col, d_type in df.dtypes.items():
+            if pd.api.extensions.ExtensionDtype.is_dtype(d_type):
+                continue
+            if not np.issubdtype(d_type, np.number):
+                df[col] = df[col].astype("category")
+                enable_categorical = True
+
         if not features:
-            return xgb.DMatrix(df)
+            return xgb.DMatrix(df, enable_categorical=enable_categorical)
         else:
             feature_names = []
             feature_types = []
@@ -95,4 +103,9 @@ class XGBoostDMatrixHandler(base_handler.BaseDataHandler["xgboost.DMatrix"]):
                 assert isinstance(feature, core.FeatureSpec), "Invalid feature kind."
                 feature_names.append(feature.name)
                 feature_types.append(feature._dtype._numpy_type)
-            return xgb.DMatrix(df, feature_names=feature_names, feature_types=feature_types)
+            return xgb.DMatrix(
+                df,
+                feature_names=feature_names,
+                feature_types=feature_types,
+                enable_categorical=enable_categorical,
+            )