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

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

@@ -1,7 +1,7 @@
 import inspect
 import logging
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -24,7 +24,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def _validate_sentence_transformers_signatures(sigs: Dict[str, model_signature.ModelSignature]) -> None:
+def _validate_sentence_transformers_signatures(sigs: dict[str, model_signature.ModelSignature]) -> None:
     if list(sigs.keys()) != ["encode"]:
         raise ValueError("target_methods can only be ['encode']")
 
@@ -48,7 +48,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
     HANDLER_TYPE = "sentence_transformers"
     HANDLER_VERSION = "2024-03-15"
     _MIN_SNOWPARK_ML_VERSION = "1.3.1"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
     DEFAULT_TARGET_METHODS = ["encode"]
@@ -166,7 +166,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
             ],
             check_local_version=True,
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @staticmethod
     def _get_device_config(**kwargs: Unpack[model_types.SentenceTransformersLoadOptions]) -> Optional[str]:
@@ -224,7 +224,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         def _create_custom_model(
             raw_model: "sentence_transformers.SentenceTransformer",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             batch_size = cast(
                 model_meta_schema.SentenceTransformersModelBlobOptions, model_meta.models[model_meta.name].options
             ).get("batch_size", None)

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

@@ -1,13 +1,13 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, Union, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, Sequence, Union, cast, final
 
 import cloudpickle
 import numpy as np
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
-from snowflake.ml._internal import type_utils
+from snowflake.ml._internal import env, type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
@@ -19,7 +19,6 @@ 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
@@ -39,6 +38,35 @@ def _unpack_container_runtime_pipeline(model: "sklearn.pipeline.Pipeline") -> "s
     return model
 
 
+def _apply_transforms_up_to_last_step(
+    model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
+    data: model_types.SupportedDataType,
+    input_feature_names: Optional[list[str]] = None,
+) -> pd.DataFrame:
+    """Apply all transformations in the sklearn pipeline model up to the last step."""
+    transformed_data = data
+    output_features_names = input_feature_names
+
+    if type_utils.LazyType("sklearn.pipeline.Pipeline").isinstance(model):
+        for step_name, step in model.steps[:-1]:  # type: ignore[attr-defined]
+            if not hasattr(step, "transform"):
+                raise ValueError(f"Step '{step_name}' does not have a 'transform' method.")
+            transformed_data = step.transform(transformed_data)
+            if output_features_names is None:
+                continue
+            elif hasattr(step, "get_feature_names_out"):
+                output_features_names = step.get_feature_names_out(output_features_names)
+            else:
+                raise ValueError(
+                    f"Step '{step_name}' in the pipeline does not have a 'get_feature_names_out' method. "
+                    "Feature names cannot be propagated."
+                )
+    if type_utils.LazyType("scipy.sparse.csr_matrix").isinstance(transformed_data):
+        # Convert to dense array if it's a sparse matrix
+        transformed_data = transformed_data.toarray()  # type: ignore[attr-defined]
+    return pd.DataFrame(transformed_data, columns=output_features_names)
+
+
 @final
 class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]]):
     """Handler for scikit-learn based model.
@@ -49,7 +77,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
     HANDLER_TYPE = "sklearn"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     DEFAULT_TARGET_METHODS = [
         "predict",
@@ -59,7 +87,9 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         "decision_function",
         "score_samples",
     ]
-    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
+
+    # Prioritize predict_proba as it gives multi-class probabilities
+    EXPLAIN_TARGET_METHODS = ["predict_proba", "predict", "predict_log_proba"]
 
     @classmethod
     def can_handle(
@@ -113,7 +143,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 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):
+        if env.IN_ML_RUNTIME and isinstance(model, sklearn.pipeline.Pipeline):
             model = _unpack_container_runtime_pipeline(model)
 
         if not is_sub_model:
@@ -161,17 +191,38 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                         stacklevel=1,
                     )
                     enable_explainability = False
-                elif model_meta.task == model_types.Task.UNKNOWN or explain_target_method is None:
+                elif model_meta.task == model_types.Task.UNKNOWN:
+                    enable_explainability = False
+                elif explain_target_method is None:
                     enable_explainability = False
                 else:
                     enable_explainability = True
             if enable_explainability:
-                model_meta = handlers_utils.add_explain_method_signature(
-                    model_meta=model_meta,
-                    explain_method="explain",
-                    target_method=explain_target_method,
-                    output_return_type=model_task_and_output_type.output_type,
+                explain_target_method = str(explain_target_method)  # mypy complains if we don't cast to str here
+
+                input_signature = handlers_utils.get_input_signature(model_meta, explain_target_method)
+                transformed_background_data = _apply_transforms_up_to_last_step(
+                    model=model,
+                    data=background_data,
+                    input_feature_names=[spec.name for spec in input_signature],
                 )
+
+                try:
+                    model_meta = handlers_utils.add_inferred_explain_method_signature(
+                        model_meta=model_meta,
+                        explain_method="explain",
+                        target_method=explain_target_method,
+                        background_data=background_data,
+                        explain_fn=cls._build_explain_fn(model, background_data, input_signature),
+                        output_feature_names=transformed_background_data.columns,
+                    )
+                except ValueError:
+                    if kwargs.get("enable_explainability", None):
+                        # user explicitly enabled explainability, so we should raise the error
+                        raise ValueError(
+                            "Explainability for this model is not supported. Please set `enable_explainability=False`"
+                        )
+
                 handlers_utils.save_background_data(
                     model_blobs_dir_path,
                     cls.EXPLAIN_ARTIFACTS_DIR,
@@ -223,11 +274,13 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                     )
 
         if enable_explainability:
-            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
+            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap>=0.46.0", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
 
         model_meta.env.include_if_absent(
-            [model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn")],
+            [
+                model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn"),
+            ],
             check_local_version=True,
         )
 
@@ -265,7 +318,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         def _create_custom_model(
             raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
                 signature: model_signature.ModelSignature,
@@ -287,37 +340,8 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
 
                 @custom_model.inference_api
                 def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
-                    import shap
-
-                    try:
-                        explainer = shap.Explainer(raw_model, background_data)
-                        df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer(X).values)
-                    except TypeError:
-                        try:
-                            dtype_map = {spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in signature.inputs}
-
-                            if isinstance(X, pd.DataFrame):
-                                X = X.astype(dtype_map, copy=False)
-                            if hasattr(raw_model, "predict_proba"):
-                                if isinstance(X, np.ndarray):
-                                    explanations = shap.Explainer(
-                                        raw_model.predict_proba, background_data.values  # type: ignore[union-attr]
-                                    )(X).values
-                                else:
-                                    explanations = shap.Explainer(raw_model.predict_proba, background_data)(X).values
-                            elif hasattr(raw_model, "predict"):
-                                if isinstance(X, np.ndarray):
-                                    explanations = shap.Explainer(
-                                        raw_model.predict, background_data.values  # type: ignore[union-attr]
-                                    )(X).values
-                                else:
-                                    explanations = shap.Explainer(raw_model.predict, background_data)(X).values
-                            else:
-                                raise ValueError("Missing any supported target method to explain.")
-                            df = handlers_utils.convert_explanations_to_2D_df(raw_model, explanations)
-                        except TypeError as e:
-                            raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
-                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+                    fn = cls._build_explain_fn(raw_model, background_data, signature.inputs)
+                    return model_signature_utils.rename_pandas_df(fn(X), signature.outputs)
 
                 if target_method == "explain":
                     return explain_fn
@@ -340,3 +364,37 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         skl_model = _SKLModel(custom_model.ModelContext())
 
         return skl_model
+
+    @classmethod
+    def _build_explain_fn(
+        cls,
+        model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
+        background_data: model_types.SupportedDataType,
+        input_specs: Sequence[model_signature.BaseFeatureSpec],
+    ) -> Callable[[model_types.SupportedDataType], pd.DataFrame]:
+        import shap
+        import sklearn.pipeline
+
+        transformed_bg_data = _apply_transforms_up_to_last_step(model, background_data)
+
+        def explain_fn(data: model_types.SupportedDataType) -> pd.DataFrame:
+            transformed_data = _apply_transforms_up_to_last_step(model, data)
+            predictor = model[-1] if isinstance(model, sklearn.pipeline.Pipeline) else model
+            try:
+                explainer = shap.Explainer(predictor, transformed_bg_data)
+                return handlers_utils.convert_explanations_to_2D_df(model, explainer(transformed_data).values)
+            except TypeError:
+                if isinstance(data, pd.DataFrame):
+                    dtype_map = {spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in input_specs}
+                    transformed_data = _apply_transforms_up_to_last_step(model, data.astype(dtype_map))
+                for explain_target_method in cls.EXPLAIN_TARGET_METHODS:
+                    if not hasattr(predictor, explain_target_method):
+                        continue
+                    explain_target_method_fn = getattr(predictor, explain_target_method)
+                    explanations = shap.Explainer(explain_target_method_fn, transformed_bg_data.values)(
+                        transformed_data.to_numpy()
+                    ).values
+                    return handlers_utils.convert_explanations_to_2D_df(model, explanations)
+                raise ValueError("Missing any supported target method to explain.")
+
+        return explain_fn

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

@@ -1,6 +1,6 @@
 import os
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast, final
 
 import cloudpickle
 import numpy as np
@@ -36,7 +36,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     HANDLER_TYPE = "snowml"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     DEFAULT_TARGET_METHODS = ["predict", "transform", "predict_proba", "predict_log_proba", "decision_function"]
     EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
@@ -264,7 +264,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         def _create_custom_model(
             raw_model: "BaseEstimator",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "BaseEstimator",
                 signature: model_signature.ModelSignature,

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

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from packaging import version
@@ -38,7 +38,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
     HANDLER_TYPE = "tensorflow"
     HANDLER_VERSION = "2025-03-01"
     _MIN_SNOWPARK_ML_VERSION = "1.8.0"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {
         "2023-12-01": tensorflow_migrator_2023_12_01.TensorflowHandlerMigrator20231201,
         "2025-01-01": tensorflow_migrator_2025_01_01.TensorflowHandlerMigrator20250101,
     }
@@ -88,6 +88,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         import tensorflow
 
         assert isinstance(model, tensorflow.Module)
+        multiple_inputs = kwargs.get("multiple_inputs", False)
 
         is_keras_model = type_utils.LazyType("keras.Model").isinstance(model)
         is_tf_keras_model = type_utils.LazyType("tf_keras.Model").isinstance(model)
@@ -112,8 +113,6 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
                 default_target_methods=default_target_methods,
             )
 
-            multiple_inputs = kwargs.get("multiple_inputs", False)
-
             if is_keras_model and len(target_methods) > 1:
                 raise ValueError("Keras model can only have one target method.")
 
@@ -188,7 +187,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             dependencies,
             check_local_version=True,
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -198,7 +197,6 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blobs_dir_path: str,
         **kwargs: Unpack[model_types.TensorflowLoadOptions],
     ) -> "tensorflow.Module":
-        os.environ["TF_USE_LEGACY_KERAS"] = "1"
         import tensorflow
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
@@ -209,7 +207,12 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         load_path = os.path.join(model_blob_path, model_blob_filename)
         save_format = model_blob_options.get("save_format", "keras_tf")
         if save_format == "keras_tf":
-            m = tensorflow.keras.models.load_model(load_path)
+            if version.parse(tensorflow.keras.__version__) >= version.parse("3.0.0"):
+                import tf_keras
+
+                m = tf_keras.models.load_model(load_path)
+            else:
+                m = tensorflow.keras.models.load_model(load_path)
         else:
             m = tensorflow.saved_model.load(load_path)
 
@@ -230,7 +233,7 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         def _create_custom_model(
             raw_model: "tensorflow.Module",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             multiple_inputs = cast(
                 model_meta_schema.TensorflowModelBlobOptions, model_meta.models[model_meta.name].options
             )["multiple_inputs"]

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

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
+from typing import TYPE_CHECKING, Callable, Optional, cast, final
 
 import pandas as pd
 from typing_extensions import TypeGuard, Unpack
@@ -36,7 +36,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     HANDLER_TYPE = "torchscript"
     HANDLER_VERSION = "2025-03-01"
     _MIN_SNOWPARK_ML_VERSION = "1.8.0"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {
         "2023-12-01": torchscript_migrator_2023_12_01.TorchScriptHandlerMigrator20231201
     }
 
@@ -76,6 +76,8 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         if enable_explainability:
             raise NotImplementedError("Explainability is not supported for Torch Script model.")
 
+        multiple_inputs = kwargs.get("multiple_inputs", False)
+
         import torch
 
         assert isinstance(model, torch.jit.ScriptModule)
@@ -87,8 +89,6 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
                 default_target_methods=cls.DEFAULT_TARGET_METHODS,
             )
 
-            multiple_inputs = kwargs.get("multiple_inputs", False)
-
             def get_prediction(
                 target_method_name: str, sample_input_data: "model_types.SupportedLocalDataType"
             ) -> model_types.SupportedLocalDataType:
@@ -141,7 +141,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         model_meta.env.include_if_absent(
             [model_env.ModelDependency(requirement="pytorch", pip_name="torch")], check_local_version=True
         )
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -181,7 +181,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
         def _create_custom_model(
             raw_model: "torch.jit.ScriptModule",
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "torch.jit.ScriptModule",
                 signature: model_signature.ModelSignature,

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

@@ -1,17 +1,7 @@
 # mypy: disable-error-code="import"
 import os
 import warnings
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    Optional,
-    Type,
-    Union,
-    cast,
-    final,
-)
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
 import numpy as np
 import pandas as pd
@@ -44,7 +34,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
     HANDLER_TYPE = "xgboost"
     HANDLER_VERSION = "2023-12-01"
     _MIN_SNOWPARK_ML_VERSION = "1.0.12"
-    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+    _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model.ubj"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
@@ -154,7 +144,12 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
-            options=model_meta_schema.XgboostModelBlobOptions({"xgb_estimator_type": model.__class__.__name__}),
+            options=model_meta_schema.XgboostModelBlobOptions(
+                {
+                    "xgb_estimator_type": model.__class__.__name__,
+                    "enable_categorical": getattr(model, "enable_categorical", False),
+                }
+            ),
         )
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
@@ -162,11 +157,6 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         model_meta.env.include_if_absent(
             [
                 model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn"),
-            ],
-            check_local_version=True,
-        )
-        model_meta.env.include_if_absent(
-            [
                 model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
             ],
             check_local_version=True,
@@ -175,7 +165,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap>=0.46.0", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
-        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+        model_meta.env.cuda_version = kwargs.get("cuda_version", handlers_utils.get_default_cuda_version())
 
     @classmethod
     def load_model(
@@ -200,6 +190,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
             raise ValueError("Type of XGB estimator is illegal.")
         m = getattr(xgboost, xgb_estimator_type)()
         m.load_model(os.path.join(model_blob_path, model_blob_filename))
+        m.enable_categorical = model_blob_options.get("enable_categorical", False)
 
         if kwargs.get("use_gpu", False):
             assert type(kwargs.get("use_gpu", False)) == bool
@@ -227,7 +218,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
         def _create_custom_model(
             raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
             model_meta: model_meta_api.ModelMetadata,
-        ) -> Type[custom_model.CustomModel]:
+        ) -> type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["xgboost.Booster", "xgboost.XGBModel"],
                 signature: model_signature.ModelSignature,
@@ -235,8 +226,16 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
                 def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+                    enable_categorical = False
+                    for col, d_type in X.dtypes.items():
+                        if pd.api.extensions.ExtensionDtype.is_dtype(d_type):
+                            continue
+                        if not np.issubdtype(d_type, np.number):
+                            # categorical columns are converted to numpy's str dtype
+                            X[col] = X[col].astype("category")
+                            enable_categorical = True
                     if isinstance(raw_model, xgboost.Booster):
-                        X = xgboost.DMatrix(X)
+                        X = xgboost.DMatrix(X, enable_categorical=enable_categorical)
 
                     res = getattr(raw_model, target_method)(X)
 
@@ -261,7 +260,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     return explain_fn
                 return fn
 
-            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            type_method_dict: dict[str, Any] = {"_raw_model": raw_model}
             for target_method_name, sig in model_meta.signatures.items():
                 type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 

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

@@ -1,4 +1,4 @@
-from typing import Dict, cast
+from typing import cast
 
 from typing_extensions import Unpack
 
@@ -25,7 +25,7 @@ class ModelBlobMeta:
         self.handler_version = kwargs["handler_version"]
         self.function_properties = kwargs.get("function_properties", {})
 
-        self.artifacts: Dict[str, str] = {}
+        self.artifacts: dict[str, str] = {}
         artifacts = kwargs.get("artifacts", None)
         if artifacts:
             self.artifacts = artifacts