snowflake-ml-python 1.31.0__tar.gz → 1.32.0__tar.gz

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/CHANGELOG.md

@@ -1,6 +1,29 @@
 # Release History
 
-## 1.31.0
+## 1.32.0
+
+### New Features
+
+* Registry: Extended `ParamSpec` support to scikit-learn models. Parameters declared in the model signature
+  can now be passed at inference time, matching the existing support for custom models.
+* Registry: `infer_signature()` now accepts a `dict` for the `params` argument (e.g.,
+  `params={"temperature": 0.7, "max_tokens": 100}`), automatically inferring `ParamSpec` objects
+  from the Python value types. The existing `Sequence[BaseParamSpec]` form is still supported.
+
+### Bug Fixes
+
+* Experiment Tracking live logging (PrPr): Fixed a bug where the tracebacks for uncaught exceptions were not logged
+  when the code was running inside Snowflake notebooks.
+
+### Behavior Changes
+
+* Registry: The models from `huggingface.TransformersPipeline` and `transformers.Pipeline` will default to safetensors
+  file format by default by using `safe_serialization=True` parameter to the `save_pretrained` call. This creates
+  safetensors files instead of PyTorch binaries or pt files, addressing security vulnerability CVE-2025-32434.
+
+### Deprecations
+
+## 1.31.0 (2026-03-19)
 
 ### New Features
 
@@ -18,7 +41,7 @@
 * Registry: Fixed a bug where invalid parameter types (e.g., passing a string to an integer parameter)
   were silently coerced instead of raising an error. Parameters are now strictly validated against their
   declared types before inference. This may break code that relied on the silent coercion behavior
-  (e.g., `model.run(max_tokens="100")` or `model.run(some_int_param=True)`).
+  (e.g., `mv.run(..., params={"max_tokens": "100"})` or `mv.run(..., params={"some_int_param": True})`).
 * ML Job: Fixed explicitly named `MLJobDefinition` objects so they can be overwritten in place and invoked
   repeatedly without job name collisions; the `generate_suffix` argument has been removed.
 * Experiment Tracking live logging (PrPr): Fixed a bug that generated too many lines of log that consist entirely of

{snowflake_ml_python-1.31.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.31.0
+Version: 1.32.0
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -416,7 +416,30 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
-## 1.31.0
+## 1.32.0
+
+### New Features
+
+* Registry: Extended `ParamSpec` support to scikit-learn models. Parameters declared in the model signature
+  can now be passed at inference time, matching the existing support for custom models.
+* Registry: `infer_signature()` now accepts a `dict` for the `params` argument (e.g.,
+  `params={"temperature": 0.7, "max_tokens": 100}`), automatically inferring `ParamSpec` objects
+  from the Python value types. The existing `Sequence[BaseParamSpec]` form is still supported.
+
+### Bug Fixes
+
+* Experiment Tracking live logging (PrPr): Fixed a bug where the tracebacks for uncaught exceptions were not logged
+  when the code was running inside Snowflake notebooks.
+
+### Behavior Changes
+
+* Registry: The models from `huggingface.TransformersPipeline` and `transformers.Pipeline` will default to safetensors
+  file format by default by using `safe_serialization=True` parameter to the `save_pretrained` call. This creates
+  safetensors files instead of PyTorch binaries or pt files, addressing security vulnerability CVE-2025-32434.
+
+### Deprecations
+
+## 1.31.0 (2026-03-19)
 
 ### New Features
 
@@ -434,7 +457,7 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 * Registry: Fixed a bug where invalid parameter types (e.g., passing a string to an integer parameter)
   were silently coerced instead of raising an error. Parameters are now strictly validated against their
   declared types before inference. This may break code that relied on the silent coercion behavior
-  (e.g., `model.run(max_tokens="100")` or `model.run(some_int_param=True)`).
+  (e.g., `mv.run(..., params={"max_tokens": "100"})` or `mv.run(..., params={"some_int_param": True})`).
 * ML Job: Fixed explicitly named `MLJobDefinition` objects so they can be overwritten in place and invoked
   repeatedly without job name collisions; the `generate_suffix` argument has been removed.
 * Experiment Tracking live logging (PrPr): Fixed a bug that generated too many lines of log that consist entirely of

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/experiment/_logging/experiment_logging_context.py

@@ -1,5 +1,6 @@
 import contextlib
 from dataclasses import dataclass
+from typing import Any, Callable, Optional
 
 from snowflake.ml._internal.utils import tee
 from snowflake.ml.experiment._logging.experiment_logger import ExperimentLogger
@@ -11,3 +12,5 @@ class ExperimentLoggingContext:
     stderr_logger: ExperimentLogger
     stdout_ctx: contextlib.redirect_stdout[tee.OutputTee]
     stderr_ctx: contextlib.redirect_stderr[tee.OutputTee]
+    ipython_instance: Optional[Any]
+    original_showtraceback: Optional[Callable[..., None]]

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/experiment/experiment_tracking.py

@@ -3,6 +3,7 @@ import functools
 import json
 import sys
 import warnings
+from collections.abc import Callable
 from typing import Any, Optional, Union
 from urllib.parse import quote
 
@@ -617,6 +618,39 @@ class ExperimentTracking:
         self._run = None
         self._unpatch_stdout_and_stderr()
 
+    def _try_patch_ipython_showtraceback(
+        self, stderr_logger: experiment_logging.ExperimentLogger
+    ) -> tuple[Any, Optional[Callable[..., None]]]:
+        try:
+            from IPython import get_ipython
+
+            ip = get_ipython()
+            if ip:
+                original_showtraceback = ip._showtraceback
+
+                def _patched_showtraceback(
+                    etype: type[BaseException],
+                    evalue: Optional[BaseException],
+                    stb: list[str],
+                ) -> None:
+                    try:
+                        text = ip.InteractiveTB.stb2text(stb)
+                        stderr_logger.write(text + "\n")
+                        stderr_logger.flush()
+                    except Exception:
+                        pass
+                    if original_showtraceback is not None:  # for mypy
+                        original_showtraceback(etype, evalue, stb)
+
+                ip._showtraceback = _patched_showtraceback
+
+                return ip, original_showtraceback
+
+        except Exception:
+            pass  # Never let IPython patching break logging setup
+
+        return None, None
+
     def _patch_stdout_and_stderr(self) -> None:
         if not self._live_logging_enabled:
             return
@@ -626,6 +660,8 @@ class ExperimentTracking:
 
         assert self._experiment is not None and self._run is not None  # for mypy
         stdout_logger, stderr_logger, stdout_ctx, stderr_ctx = None, None, None, None  # for exception handling
+        ip: Any = None
+        original_showtraceback: Optional[Callable[..., None]] = None
         try:
             exp_id = self._sql_client.get_experiment_id(self._experiment.name)
             run_id = self._sql_client.get_run_id(
@@ -636,15 +672,23 @@ class ExperimentTracking:
             stderr_logger = experiment_logging.ExperimentLogger(exp_id=exp_id, run_id=run_id, stream="STDERR")
             stdout_ctx = contextlib.redirect_stdout(tee.OutputTee(sys.stdout, stdout_logger))
             stderr_ctx = contextlib.redirect_stderr(tee.OutputTee(sys.stderr, stderr_logger))
+            stdout_ctx.__enter__()
+            stderr_ctx.__enter__()
+
+            ip, original_showtraceback = self._try_patch_ipython_showtraceback(stderr_logger)
+
             self._logging_context = experiment_logging.ExperimentLoggingContext(
                 stdout_logger=stdout_logger,
                 stderr_logger=stderr_logger,
                 stdout_ctx=stdout_ctx,
                 stderr_ctx=stderr_ctx,
+                ipython_instance=ip,
+                original_showtraceback=original_showtraceback,
             )
-            stdout_ctx.__enter__()
-            stderr_ctx.__enter__()
+
         except Exception as e:
+            if ip:
+                ip._showtraceback = original_showtraceback
             if stdout_ctx:
                 stdout_ctx.__exit__(None, None, None)
             if stderr_ctx:
@@ -661,6 +705,8 @@ class ExperimentTracking:
 
     def _unpatch_stdout_and_stderr(self) -> None:
         if self._logging_context:
+            if self._logging_context.ipython_instance:
+                self._logging_context.ipython_instance._showtraceback = self._logging_context.original_showtraceback
             self._logging_context.stdout_ctx.__exit__(None, None, None)
             self._logging_context.stderr_ctx.__exit__(None, None, None)
             self._logging_context.stdout_logger.close()

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/feature_store/feature_store.py

@@ -2772,7 +2772,7 @@ class FeatureStore:
 
         # Build SELECT and WHERE clauses using helper methods
         select_clause = self._build_select_clause_and_validate(feature_view, feature_names, include_join_keys=True)
-        where_clause = self._build_where_clause_for_keys(feature_view, keys)
+        where_clause, _ = self._build_where_clause_for_keys(feature_view, keys)
 
         query = f"SELECT {select_clause} FROM {table_name}{where_clause}"
         return self._session.sql(query)
@@ -2804,7 +2804,7 @@ class FeatureStore:
         quoted_keys_str = ", ".join(quoted_keys)
 
         # Build WHERE clause for key filtering (if any)
-        where_clause = self._build_where_clause_for_keys(feature_view, keys)
+        where_clause, _ = self._build_where_clause_for_keys(feature_view, keys)
 
         # Step 1: Create spine CTE with unique entities + CURRENT_TIMESTAMP
         spine_cte = f"""
@@ -2850,10 +2850,27 @@ class FeatureStore:
         return self._session.sql(full_query)
 
     def _read_from_online_store(
-        self, feature_view: FeatureView, keys: Optional[list[list[str]]], feature_names: Optional[list[str]]
+        self,
+        feature_view: FeatureView,
+        keys: Optional[list[list[str]]],
+        feature_names: Optional[list[str]],
     ) -> DataFrame:
-        """Read feature values from the online store with optional key filtering."""
-        # Check if online store is enabled
+        """Read feature values from the online store with optional key filtering.
+
+        Uses bind variables for single-key lookups and literal interpolation
+        for batch lookups.
+
+        Args:
+            feature_view: The registered feature view to read from.
+            keys: Optional list of key value lists to filter by.
+            feature_names: Optional list of feature names to return.
+
+        Returns:
+            Snowpark DataFrame containing the feature view data.
+
+        Raises:
+            SnowflakeMLException: If online store is not enabled.
+        """
         if not feature_view.online:
             raise snowml_exceptions.SnowflakeMLException(
                 error_code=error_codes.INVALID_ARGUMENT,
@@ -2864,11 +2881,16 @@ class FeatureStore:
 
         fully_qualified_online_name = feature_view.fully_qualified_online_table_name()
 
-        # Build SELECT and WHERE clauses using helper methods
         select_clause = self._build_select_clause_and_validate(feature_view, feature_names, include_join_keys=True)
-        where_clause = self._build_where_clause_for_keys(feature_view, keys)
+        # NOTE: We use snowflake sql query binds for single-point lookups to ensure the best performance.
+        # We don't use binds for multi-point lookups because we don't have any evidence that it improves
+        # performance, but we may need to revisit this in the future.
+        use_binds = keys is not None and len(keys) == 1
+        where_clause, params = self._build_where_clause_for_keys(feature_view, keys, use_binds=use_binds)
 
         query = f"SELECT {select_clause} FROM {fully_qualified_online_name}{where_clause}"
+        if params:
+            return self._session.sql(query, params=params)
         return self._session.sql(query)
 
     @dispatch_decorator()
@@ -5050,28 +5072,31 @@ FROM SPINE{' '.join(join_clauses)}
             # Select all columns
             return "*"
 
-    def _build_where_clause_for_keys(self, feature_view: FeatureView, keys: Optional[list[list[str]]]) -> str:
+    def _build_where_clause_for_keys(
+        self, feature_view: FeatureView, keys: Optional[list[list[Any]]], use_binds: bool = False
+    ) -> tuple[str, list[Any]]:
         """Build WHERE clause for key filtering.
 
         Args:
-            feature_view: The feature view to build the clause for
-            keys: Optional list of key value lists to filter by
+            feature_view: The feature view to build the clause for.
+            keys: Optional list of key value lists to filter by.
+            use_binds: If True, use ``?`` placeholders and return bind params.
+                If False, interpolate values as string literals.
 
         Returns:
-            WHERE clause string (empty if no keys provided)
+            Tuple of (WHERE clause string, bind params list). Params is empty
+            when use_binds is False or no keys are provided.
 
         Raises:
-            SnowflakeMLException: If key structure is invalid
+            SnowflakeMLException: If key structure is invalid.
         """
         if not keys:
-            return ""
+            return "", []
 
-        # Get join keys from entities for key filtering
         all_join_keys = []
         for entity in feature_view.entities:
             all_join_keys.extend([key.resolved() for key in entity.join_keys])
 
-        # Validate key structure
         for key_values in keys:
             if len(key_values) != len(all_join_keys):
                 raise snowml_exceptions.SnowflakeMLException(
@@ -5082,15 +5107,20 @@ FROM SPINE{' '.join(join_clauses)}
                     ),
                 )
 
+        params: list[Any] = []
         where_conditions = []
         for key_values in keys:
             key_conditions = []
             for join_key, value in zip(all_join_keys, key_values):
-                safe_value = str(value).replace("'", "''")
-                key_conditions.append(f"\"{join_key}\" = '{safe_value}'")
+                if use_binds:
+                    key_conditions.append(f'"{join_key}" = ?')
+                    params.append(value)
+                else:
+                    safe_value = str(value).replace("'", "''")
+                    key_conditions.append(f"\"{join_key}\" = '{safe_value}'")
             where_conditions.append(f"({' AND '.join(key_conditions)})")
 
-        return f" WHERE {' OR '.join(where_conditions)}"
+        return f" WHERE {' OR '.join(where_conditions)}", params
 
     def _get_store_type(self, store_type: Union[fv_mod.StoreType, str]) -> fv_mod.StoreType:
         """Return a StoreType enum from a Union[StoreType, str].

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/jobs/job.py

@@ -287,10 +287,17 @@ class MLJob(Generic[T], SerializableSessionMixin):
             Optional[str]: The Ray dashboard URL if the job is running and has a Ray dashboard endpoint,
                 None otherwise.
         """
+        pool_info = _get_compute_pool_info(self._session, self._compute_pool)
+        if pool_info["instance_family"] == "CPU_X64_XS":
+            logger.warning(
+                "Ray dashboard is not supported on XS compute pools."
+                " Please use a larger compute pool if you want to use Ray dashboard."
+            )
+            return None
         if self.status != JOB_RUNNING_STATUS:
             logger.warning("Ray dashboard is not available for non-running jobs")
             return None
-        rows = self._session.sql(f"show endpoints in service {self.id}").collect()
+        rows = query_helper.run_query(self._session, "SHOW ENDPOINTS IN SERVICE IDENTIFIER(?)", params=(self.id,))
         for row in rows:
             if row["name"] == RAY_DASHBOARD_ENDPOINT_NAME:
                 ingress_url = row["ingress_url"]

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/model/_client/model/model_version_impl.py

@@ -20,7 +20,7 @@ from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._model_composer.model_method import utils as model_method_utils
-from snowflake.ml.model._packager.model_handlers import snowmlmodel
+from snowflake.ml.model._packager.model_handlers import huggingface, snowmlmodel
 from snowflake.ml.model._packager.model_meta import model_meta_schema
 from snowflake.ml.model._signatures import core
 from snowflake.snowpark import Session, async_job, dataframe
@@ -779,6 +779,14 @@ class ModelVersion(lineage_node.LineageNode):
         column_handling = input_spec.column_handling
         partition_columns = [input_spec.partition_column] if input_spec.partition_column is not None else None
 
+        if partition_columns is not None:
+            model_spec = self._get_model_spec(statement_params)
+            if model_spec.get("model_type") == huggingface.TransformersPipelineHandler.HANDLER_TYPE:
+                raise ValueError(
+                    "partition_column is not supported for HuggingFace pipeline models in batch inference jobs. "
+                    "Please remove the partition_column from InputSpec."
+                )
+
         if job_spec is None:
             job_spec = batch_inference_specs.JobSpec()
 
@@ -817,6 +825,16 @@ class ModelVersion(lineage_node.LineageNode):
 
         target_function_info = self._get_function_info(function_name=job_spec.function_name)
 
+        if (
+            partition_columns is not None
+            and target_function_info["target_method_function_type"]
+            == model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+        ):
+            raise ValueError(
+                "partition_column is not supported for FUNCTION type methods in batch inference jobs. "
+                "Only TABLE_FUNCTION type methods support partitioning."
+            )
+
         return self._service_ops.invoke_batch_job_method(
             # model version info
             model_name=self._model_name,

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/model/_packager/model_env/model_env.py

@@ -243,14 +243,17 @@ class ModelEnv:
                 channel, spec = spec_conda
                 self._conda_dependencies[channel].remove(spec)
 
-    def generate_env_for_cuda(self) -> None:
+    def _is_pip_only_path(self) -> bool:
+        """Check if this model will use the pip-only packaging path."""
+        has_conda_deps = any(len(deps) > 0 for deps in self._conda_dependencies.values())
+        return self.prefer_pip and _ENABLE_PIP_ONLY_PACKAGING and not has_conda_deps
 
-        # Insert py-xgboost-gpu only for XGBoost versions < 3.0.0
+    def _substitute_conda_gpu_packages(self) -> None:
+        """Replace conda packages with their GPU variants (xgboost→py-xgboost-gpu, tensorflow→tensorflow-gpu)."""
         xgboost_spec = env_utils.find_dep_spec(
             self._conda_dependencies, self._pip_requirements, conda_pkg_name="xgboost", remove_spec=False
         )
         if xgboost_spec:
-            # Only handle explicitly pinned versions. Insert GPU variant iff pinned major < 3.
             pinned_major: Optional[int] = None
             for spec in xgboost_spec.specifier:
                 if spec.operator in ("==", "===", ">", ">="):
@@ -279,6 +282,15 @@ class ModelEnv:
                 check_local_version=False,
             )
 
+    def generate_env_for_cuda(self) -> None:
+        # Conda-specific GPU substitutions only apply on the conda path, not the pip path.
+        include_conda_gpu_substitutions = not self._is_pip_only_path()
+
+        if include_conda_gpu_substitutions:
+            self._substitute_conda_gpu_packages()
+
+        # TODO: Implicit CAI for PyTorch (auto-add torch when needed) can be added on top of this flow later.
+
         transformers_spec = env_utils.find_dep_spec(
             self._conda_dependencies, self._pip_requirements, conda_pkg_name="transformers", remove_spec=False
         )

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/model/_packager/model_handlers/_utils.py

@@ -136,6 +136,10 @@ def add_inferred_explain_method_signature(
         output_feature_names=suffixed_output_names,
     )
 
+    target_sig = model_meta.signatures.get(target_method)
+    if target_sig is not None and target_sig.params is not None:
+        sig = model_signature.ModelSignature(inputs=sig.inputs, outputs=sig.outputs, params=target_sig.params)
+
     model_meta.signatures[explain_method] = sig
     return model_meta
 

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/model/_packager/model_handlers/huggingface.py

@@ -240,7 +240,8 @@ class TransformersPipelineHandler(
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             model.save_pretrained(  # type:ignore[attr-defined]
-                save_path
+                save_path,
+                safe_serialization=True,  # creates safetensors instead of pytorch binaries or pt files
             )
             handlers_utils.save_transformers_config_with_auto_map(
                 save_path,
@@ -553,6 +554,31 @@ class TransformersPipelineHandler(
                         input_col = signature.inputs[0].name
                         images = [Image.open(io.BytesIO(img_bytes)) for img_bytes in X[input_col].to_list()]
                         temp_res = getattr(raw_model, target_method)(images)
+                    elif isinstance(
+                        raw_model,
+                        (
+                            transformers.ImageToTextPipeline,
+                            transformers.ImageFeatureExtractionPipeline,
+                            transformers.ObjectDetectionPipeline,
+                        ),
+                    ):
+                        # Image pipelines that need bytes→PIL conversion.
+                        # HuggingFace's load_image() does not accept raw bytes.
+                        from PIL import Image
+
+                        input_col = signature.inputs[0].name
+                        if len(signature.inputs) > 1:
+                            # Multi-input: convert image bytes to PIL, pass other columns as-is
+                            def process_image_row(row: pd.Series) -> Any:
+                                pil_image = Image.open(io.BytesIO(row[input_col]))
+                                kwargs = {k: row[k] for k in row.index if k != input_col}
+                                return getattr(raw_model, target_method)(pil_image, **kwargs)
+
+                            temp_res = X.apply(process_image_row, axis=1).to_list()
+                        else:
+                            # Single-input: convert all image bytes to PIL
+                            images = [Image.open(io.BytesIO(img_bytes)) for img_bytes in X[input_col].to_list()]
+                            temp_res = getattr(raw_model, target_method)(images)
                     elif isinstance(raw_model, transformers.AutomaticSpeechRecognitionPipeline):
                         # ASR pipeline accepts a single audio input (bytes, str, np.ndarray, or dict),
                         # not a list. Process each audio input individually.
@@ -624,6 +650,7 @@ class TransformersPipelineHandler(
                                 (
                                     transformers.FillMaskPipeline,
                                     transformers.QuestionAnsweringPipeline,
+                                    transformers.ObjectDetectionPipeline,
                                 ),
                             )
                             and X.shape[0] == 1

{snowflake_ml_python-1.31.0 → snowflake_ml_python-1.32.0}/snowflake/ml/model/_packager/model_handlers/sklearn.py

@@ -1,8 +1,9 @@
+import functools
 import logging
 import os
 import warnings
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Callable, Optional, Union, cast, final
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
 import cloudpickle
 import numpy as np
@@ -331,9 +332,9 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 target_method: str,
                 background_data: Optional[pd.DataFrame],
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
-                @custom_model.inference_api
-                def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
-                    res = getattr(raw_model, target_method)(X)
+                @custom_model._internal_inference_api
+                def fn(self: custom_model.CustomModel, X: pd.DataFrame, **method_kwargs: Any) -> pd.DataFrame:
+                    res = getattr(raw_model, target_method)(X, **method_kwargs)
 
                     if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
                         # In case of multi-output estimators, predict_proba(), decision_function(), etc., functions
@@ -344,24 +345,24 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
 
                     return model_signature_utils.rename_pandas_df(df, signature.outputs)
 
-                @custom_model.inference_api
-                def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
-                    fn = cls._build_explain_fn(raw_model, background_data, signature.inputs)
-                    return model_signature_utils.rename_pandas_df(fn(X), signature.outputs)
+                @custom_model._internal_inference_api
+                def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame, **method_kwargs: Any) -> pd.DataFrame:
+                    explain = cls._build_explain_fn(raw_model, background_data, signature.inputs, **method_kwargs)
+                    return model_signature_utils.rename_pandas_df(explain(X), signature.outputs)
 
                 if target_method == "explain":
                     return explain_fn
 
                 return fn
 
-            type_method_dict = {}
+            sklearn_type_dict: dict[str, Any] = {"_allows_kwargs": True}
             for target_method_name, sig in model_meta.signatures.items():
-                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name, background_data)
+                sklearn_type_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name, background_data)
 
             _SKLModel = type(
                 "_SKLModel",
                 (custom_model.CustomModel,),
-                type_method_dict,
+                sklearn_type_dict,
             )
 
             return _SKLModel
@@ -377,17 +378,23 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
         background_data: model_types.SupportedDataType,
         input_specs: Sequence[model_signature.BaseFeatureSpec],
+        **kwargs: dict[str, Any],
     ) -> Callable[[model_types.SupportedDataType], pd.DataFrame]:
         import shap
         import sklearn.pipeline
 
         transformed_bg_data = _apply_transforms_up_to_last_step(model, background_data)
+        predictor = model[-1] if isinstance(model, sklearn.pipeline.Pipeline) else model
+
+        if kwargs:
+            predict_with_params = functools.partial(predictor.predict, **kwargs)
+        else:
+            predict_with_params = predictor
 
         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)
+                explainer = shap.Explainer(predict_with_params, transformed_bg_data)
                 return handlers_utils.convert_explanations_to_2D_df(model, explainer(transformed_data).values).astype(
                     np.float64, errors="ignore"
                 )
@@ -398,8 +405,11 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 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)(
+                    if kwargs:
+                        target_fn = functools.partial(getattr(predictor, explain_target_method), **kwargs)
+                    else:
+                        target_fn = getattr(predictor, explain_target_method)
+                    explanations = shap.Explainer(target_fn, transformed_bg_data.values)(
                         transformed_data.to_numpy()
                     ).values
                     return handlers_utils.convert_explanations_to_2D_df(model, explanations)