snowflake-ml-python 1.21.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

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

@@ -23,24 +23,55 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+# Allowlist of supported target methods for SentenceTransformer models.
+_ALLOWED_TARGET_METHODS = ["encode", "encode_queries", "encode_documents"]
+
 
 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']")
+    """Validate signatures for SentenceTransformer models.
+
+    Args:
+        sigs: Dictionary mapping method names to their signatures.
+
+    Raises:
+        ValueError: If signatures are empty, contain unsupported methods, or violate per-method constraints.
+    """
+    # Check that signatures are non-empty
+    if not sigs:
+        raise ValueError("At least one signature must be provided.")
+
+    # Check that all methods are in the allowlist
+    unsupported_methods = set(sigs.keys()) - set(_ALLOWED_TARGET_METHODS)
+    if unsupported_methods:
+        raise ValueError(
+            f"Unsupported target methods: {sorted(unsupported_methods)}. "
+            f"Supported methods are: {_ALLOWED_TARGET_METHODS}."
+        )
 
-    if len(sigs["encode"].inputs) != 1:
-        raise ValueError("SentenceTransformer can only accept 1 input column")
+    # Validate per-method constraints
+    for method_name, sig in sigs.items():
+        if len(sig.inputs) != 1:
+            raise ValueError(f"SentenceTransformer method '{method_name}' must have exactly 1 input column.")
 
-    if len(sigs["encode"].outputs) != 1:
-        raise ValueError("SentenceTransformer can only return 1 output column")
+        if len(sig.outputs) != 1:
+            raise ValueError(f"SentenceTransformer method '{method_name}' must have exactly 1 output column.")
 
-    assert isinstance(sigs["encode"].inputs[0], model_signature.FeatureSpec)
+        # FeatureSpec is expected here; FeatureGroupSpec would indicate a nested/grouped input
+        # which SentenceTransformer does not support.
+        if not isinstance(sig.inputs[0], model_signature.FeatureSpec):
+            raise ValueError(
+                f"SentenceTransformer method '{method_name}' requires a FeatureSpec input, "
+                f"got {type(sig.inputs[0]).__name__}."
+            )
 
-    if sigs["encode"].inputs[0]._shape is not None:
-        raise ValueError("SentenceTransformer does not support input shape")
+        if sig.inputs[0]._shape is not None:
+            raise ValueError(f"SentenceTransformer method '{method_name}' does not support input shape.")
 
-    if sigs["encode"].inputs[0]._dtype != model_signature.DataType.STRING:
-        raise ValueError("SentenceTransformer only accepts string input")
+        if sig.inputs[0]._dtype != model_signature.DataType.STRING:
+            raise ValueError(
+                f"SentenceTransformer method '{method_name}' only accepts STRING input, "
+                f"got {sig.inputs[0]._dtype.name}."
+            )
 
 
 @final
@@ -51,7 +82,7 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
     _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
-    DEFAULT_TARGET_METHODS = ["encode"]
+    DEFAULT_TARGET_METHODS = ["encode", "encode_queries", "encode_documents"]
 
     @classmethod
     def can_handle(
@@ -98,8 +129,13 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
                 target_methods=kwargs.pop("target_methods", None),
                 default_target_methods=cls.DEFAULT_TARGET_METHODS,
             )
-            if target_methods != ["encode"]:
-                raise ValueError("target_methods can only be ['encode']")
+
+            # Validate target_methods
+            if not target_methods:
+                raise ValueError("At least one target method must be specified.")
+
+            if not set(target_methods).issubset(_ALLOWED_TARGET_METHODS):
+                raise ValueError(f"target_methods {target_methods} must be a subset of {_ALLOWED_TARGET_METHODS}.")
 
             def get_prediction(
                 target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
@@ -246,10 +282,10 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
 
             type_method_dict = {}
             for target_method_name, sig in model_meta.signatures.items():
-                if target_method_name == "encode":
+                if target_method_name in _ALLOWED_TARGET_METHODS:
                     type_method_dict[target_method_name] = get_prediction(raw_model, sig, target_method_name)
                 else:
-                    ValueError(f"{target_method_name} is currently not supported.")
+                    raise ValueError(f"{target_method_name} is currently not supported.")
 
             _SentenceTransformer = type(
                 "_SentenceTransformer",

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

@@ -194,7 +194,18 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
 
         if kwargs.get("use_gpu", False):
             assert type(kwargs.get("use_gpu", False)) == bool
-            gpu_params = {"tree_method": "gpu_hist", "predictor": "gpu_predictor"}
+            from packaging import version
+
+            xgb_version = version.parse(xgboost.__version__)
+            if xgb_version >= version.parse("3.1.0"):
+                # XGBoost 3.1.0+: Use device="cuda" for GPU acceleration
+                # gpu_hist and gpu_predictor were removed in XGBoost 3.1.0
+                # See: https://xgboost.readthedocs.io/en/latest/changes/v3.1.0.html
+                gpu_params = {"tree_method": "hist", "device": "cuda"}
+            else:
+                # XGBoost < 3.1.0: Use legacy gpu_hist tree_method
+                gpu_params = {"tree_method": "gpu_hist", "predictor": "gpu_predictor"}
+
             if isinstance(m, xgboost.Booster):
                 m.set_param(gpu_params)
             elif isinstance(m, xgboost.XGBModel):
@@ -256,6 +267,20 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                 @custom_model.inference_api
                 def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
                     import shap
+                    from packaging import version
+
+                    xgb_version = version.parse(xgboost.__version__)
+                    shap_version = version.parse(shap.__version__)
+
+                    # SHAP < 0.50.0 is incompatible with XGBoost >= 3.1.0 due to base_score format change
+                    # (base_score is now stored as a vector for multi-output models)
+                    # See: https://xgboost.readthedocs.io/en/latest/changes/v3.1.0.html
+                    if xgb_version >= version.parse("3.1.0") and shap_version < version.parse("0.50.0"):
+                        raise RuntimeError(
+                            f"SHAP version {shap.__version__} is incompatible with XGBoost version "
+                            f"{xgboost.__version__}. XGBoost 3.1+ changed the model format which requires "
+                            f"SHAP >= 0.50.0. Please upgrade SHAP or use XGBoost < 3.1."
+                        )
 
                     explainer = shap.TreeExplainer(raw_model)
                     df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer.shap_values(X))

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

@@ -20,6 +20,7 @@ from snowflake.ml.model._packager.model_env import model_env
 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
+from snowflake.ml.model.code_path import CodePath
 
 MODEL_METADATA_FILE = "model.yaml"
 MODEL_CODE_DIR = "code"
@@ -39,7 +40,7 @@ def create_model_metadata(
     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,
+    code_paths: Optional[list[model_types.CodePathLike]] = None,
     ext_modules: Optional[list[ModuleType]] = None,
     conda_dependencies: Optional[list[str]] = None,
     pip_requirements: Optional[list[str]] = None,
@@ -77,7 +78,8 @@ def create_model_metadata(
         **kwargs: Dict of attributes and values of the metadata. Used when loading from file.
 
     Raises:
-        ValueError: Raised when the code path contains reserved file or directory.
+        ValueError: Raised when the code path contains reserved file or directory, or destination conflicts.
+        FileNotFoundError: Raised when a code path does not exist.
 
     Yields:
         A model metadata object.
@@ -134,13 +136,44 @@ def create_model_metadata(
         os.makedirs(code_dir_path, exist_ok=True)
 
     if code_paths:
+        # Resolve all code paths and check for conflicts
+        resolved_paths: list[tuple[str, str]] = []  # (source, destination_relative)
         for code_path in code_paths:
-            # This part is to prevent users from providing code following our naming and overwrite our code.
-            if (
-                os.path.isfile(code_path) and os.path.splitext(os.path.basename(code_path))[0] == _SNOWFLAKE_PKG_NAME
-            ) or (os.path.isdir(code_path) and os.path.basename(code_path) == _SNOWFLAKE_PKG_NAME):
+            if isinstance(code_path, CodePath):
+                source, dest_relative = code_path._resolve()
+            else:
+                # String path: keep existing behavior
+                source = os.path.normpath(os.path.abspath(code_path))
+                if not os.path.exists(source):
+                    raise FileNotFoundError(f"Code path '{code_path}' does not exist (resolved to {source}).")
+                dest_relative = os.path.basename(source)
+            resolved_paths.append((source, dest_relative))
+
+        # Check for destination conflicts
+        seen: dict[str, str] = {}
+        for source, dest in resolved_paths:
+            if dest in seen:
+                raise ValueError(
+                    f"Destination path conflict: '{dest}' is targeted by both '{seen[dest]}' and '{source}'."
+                )
+            seen[dest] = source
+
+        # Copy files
+        for source, dest_relative in resolved_paths:
+            # Prevent reserved name conflicts
+            dest_name = dest_relative.split(os.sep)[0] if os.sep in dest_relative else dest_relative
+            if (os.path.isfile(source) and os.path.splitext(dest_name)[0] == _SNOWFLAKE_PKG_NAME) or (
+                os.path.isdir(source) and dest_name == _SNOWFLAKE_PKG_NAME
+            ):
                 raise ValueError("`snowflake` is a reserved name and you cannot contain that into code path.")
-            file_utils.copy_file_or_tree(code_path, code_dir_path)
+
+            parent_dir = (
+                os.path.join(code_dir_path, os.path.dirname(dest_relative))
+                if os.path.dirname(dest_relative)
+                else code_dir_path
+            )
+            os.makedirs(parent_dir, exist_ok=True)
+            file_utils.copy_file_or_tree(source, parent_dir)
 
     try:
         imported_modules = []

snowflake/ml/model/_packager/model_packager.py

@@ -49,7 +49,7 @@ class ModelPackager:
         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,
+        code_paths: Optional[list[model_types.CodePathLike]] = None,
         options: model_types.ModelSaveOption,
         task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:

snowflake/ml/model/_signatures/core.py

@@ -191,6 +191,35 @@ class DataType(Enum):
             original_exception=NotImplementedError(f"Type {snowpark_type} is not supported as a DataType."),
         )
 
+    @classmethod
+    def from_python_type(cls, python_type: type) -> "DataType":
+        """Translate Python built-in type to DataType for signature definition.
+
+        Args:
+            python_type: A Python built-in type (int, float, str, bool).
+
+        Raises:
+            SnowflakeMLException: NotImplementedError: Raised when the given Python type is not supported.
+
+        Returns:
+            Corresponding DataType.
+        """
+        python_to_snowml_type_mapping: dict[type, "DataType"] = {
+            int: DataType.INT64,
+            float: DataType.DOUBLE,
+            str: DataType.STRING,
+            bool: DataType.BOOL,
+        }
+        if python_type in python_to_snowml_type_mapping:
+            return python_to_snowml_type_mapping[python_type]
+        raise snowml_exceptions.SnowflakeMLException(
+            error_code=error_codes.NOT_IMPLEMENTED,
+            original_exception=NotImplementedError(
+                f"Python type {python_type} is not supported as a DataType. "
+                f"Supported types are: {list(python_to_snowml_type_mapping.keys())}."
+            ),
+        )
+
 
 class BaseFeatureSpec(ABC):
     """Abstract Class for specification of a feature."""
@@ -764,10 +793,17 @@ class ModelSignature:
                 the output of the model.
             params: A sequence of parameter specifications and parameter group specifications that will compose
                 the parameters of the model. Defaults to None.
+
+        Raises:
+            SnowflakeMLException: ValueError: When the parameters have duplicate names or the same
+                names as input features.
+
+        # noqa: DAR402
         """
         self._inputs = inputs
         self._outputs = outputs
         self._params = params or []
+        self._name_validation()
 
     @property
     def inputs(self) -> Sequence[BaseFeatureSpec]:
@@ -879,6 +915,55 @@ class ModelSignature:
 
         return html_utils.create_base_container("Model Signature", content)
 
+    def _name_validation(self) -> None:
+        """Validate the names of the inputs and parameters.
+
+        Names are compared case-insensitively (matches Snowflake identifier behavior).
+
+        Raises:
+            SnowflakeMLException: ValueError: When the parameters have duplicate names or the same
+                names as input features.
+        """
+        input_names: set[str] = set()
+        for input_spec in self._inputs:
+            names = (
+                [input_spec.name.upper() for spec in input_spec._specs]
+                if isinstance(input_spec, FeatureGroupSpec)
+                else [input_spec.name.upper()]
+            )
+            input_names.update(names)
+
+        param_names: set[str] = set()
+        dup_params: set[str] = set()
+        collision_names: set[str] = set()
+
+        for param in self._params:
+            names = [spec.name for spec in param.specs] if isinstance(param, ParamGroupSpec) else [param.name]
+            for name in names:
+                if name.upper() in param_names:
+                    dup_params.add(name)
+                if name.upper() in input_names:
+                    collision_names.add(name)
+                param_names.add(name.upper())
+
+        if dup_params:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(
+                    f"Found duplicate parameter named resolved as {', '.join(sorted(dup_params))}."
+                    " Parameters must have distinct names (case-insensitive)."
+                ),
+            )
+
+        if collision_names:
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.INVALID_ARGUMENT,
+                original_exception=ValueError(
+                    f"Found parameter(s) with the same name as input feature(s): {', '.join(sorted(collision_names))}."
+                    " Parameters and inputs must have distinct names (case-insensitive)."
+                ),
+            )
+
     @classmethod
     def from_mlflow_sig(cls, mlflow_sig: "mlflow.models.ModelSignature") -> "ModelSignature":
         return ModelSignature(

snowflake/ml/model/_signatures/utils.py

@@ -9,6 +9,7 @@ from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
+from snowflake.ml.model import openai_signatures
 from snowflake.ml.model._signatures import core
 
 
@@ -259,6 +260,48 @@ def huggingface_pipeline_signature_auto_infer(
             ],
         )
 
+    # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.ImageClassificationPipeline
+    if task == "image-classification":
+        return core.ModelSignature(
+            inputs=[
+                core.FeatureSpec(name="images", dtype=core.DataType.BYTES),
+            ],
+            outputs=[
+                core.FeatureGroupSpec(
+                    name="labels",
+                    specs=[
+                        core.FeatureSpec(name="label", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="score", dtype=core.DataType.DOUBLE),
+                    ],
+                    shape=(-1,),
+                ),
+            ],
+        )
+
+    # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.AutomaticSpeechRecognitionPipeline
+    if task == "automatic-speech-recognition":
+        return core.ModelSignature(
+            inputs=[
+                core.FeatureSpec(name="audio", dtype=core.DataType.BYTES),
+            ],
+            outputs=[
+                core.FeatureGroupSpec(
+                    name="outputs",
+                    specs=[
+                        core.FeatureSpec(name="text", dtype=core.DataType.STRING),
+                        core.FeatureGroupSpec(
+                            name="chunks",
+                            specs=[
+                                core.FeatureSpec(name="timestamp", dtype=core.DataType.DOUBLE, shape=(2,)),
+                                core.FeatureSpec(name="text", dtype=core.DataType.STRING),
+                            ],
+                            shape=(-1,),  # Variable length list of chunks
+                        ),
+                    ],
+                ),
+            ],
+        )
+
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.TextGenerationPipeline
     if task == "text-generation":
         if params.get("return_tensors", False):
@@ -288,6 +331,18 @@ def huggingface_pipeline_signature_auto_infer(
                 )
             ],
         )
+
+    # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.ImageTextToTextPipeline
+    if task == "image-text-to-text":
+        if params.get("return_tensors", False):
+            raise NotImplementedError(
+                f"Auto deployment for HuggingFace pipeline {task} "
+                "when `return_tensors` set to `True` has not been supported yet."
+            )
+        # Always generate a dict per input
+        # defaulting to OPENAI_CHAT_SIGNATURE_SPEC for image-text-to-text pipeline
+        return openai_signatures._OPENAI_CHAT_SIGNATURE_SPEC
+
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.Text2TextGenerationPipeline
     if task == "text2text-generation":
         if params.get("return_tensors", False):

snowflake/ml/model/code_path.py

@@ -0,0 +1,104 @@
+"""CodePath class for selective code packaging in model registry."""
+
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+_ERR_ROOT_NOT_FOUND = "CodePath: root '{root}' does not exist (resolved to {resolved})."
+_ERR_WILDCARDS_NOT_SUPPORTED = "CodePath: Wildcards are not supported in filter. Got '{filter}'. Use exact paths only."
+_ERR_FILTER_MUST_BE_RELATIVE = "CodePath: filter must be a relative path, got absolute path '{filter}'."
+_ERR_FILTER_HOME_PATH = "CodePath: filter must be a relative path, got home directory path '{filter}'."
+_ERR_FILTER_ON_FILE_ROOT = (
+    "CodePath: cannot apply filter to a file root. " "Root '{root}' is a file. Use filter only with directory roots."
+)
+_ERR_FILTER_ESCAPES_ROOT = (
+    "CodePath: filter '{filter}' escapes root directory '{root}'. " "Relative paths must stay within root."
+)
+_ERR_FILTER_NOT_FOUND = "CodePath: filter '{filter}' under root '{root}' does not exist (resolved to {resolved})."
+
+
+@dataclass(frozen=True)
+class CodePath:
+    """Specifies a code path with optional filtering for selective inclusion.
+
+    Args:
+        root: The root directory or file path (absolute or relative to cwd).
+        filter: Optional relative path under root to select a subdirectory or file.
+            The filter also determines the destination path under code/.
+
+    Examples:
+        CodePath("project/src/")                          # Copy entire src/ to code/src/
+        CodePath("project/src/", filter="utils")          # Copy utils/ to code/utils/
+        CodePath("project/src/", filter="lib/helpers")    # Copy to code/lib/helpers/
+    """
+
+    root: str
+    filter: Optional[str] = None
+
+    def __post_init__(self) -> None:
+        if self.filter == "":
+            object.__setattr__(self, "filter", None)
+
+    def __repr__(self) -> str:
+        if self.filter:
+            return f"CodePath({self.root!r}, filter={self.filter!r})"
+        return f"CodePath({self.root!r})"
+
+    def _validate_filter(self) -> Optional[str]:
+        """Validate and normalize filter, returning normalized filter or None.
+
+        Returns:
+            Normalized filter path, or None if no filter is set.
+
+        Raises:
+            ValueError: If filter contains wildcards or is an absolute path.
+        """
+        if self.filter is None:
+            return None
+
+        if any(c in self.filter for c in ["*", "?", "[", "]"]):
+            raise ValueError(_ERR_WILDCARDS_NOT_SUPPORTED.format(filter=self.filter))
+
+        if self.filter.startswith("~"):
+            raise ValueError(_ERR_FILTER_HOME_PATH.format(filter=self.filter))
+
+        filter_normalized = os.path.normpath(self.filter)
+
+        if os.path.isabs(filter_normalized):
+            raise ValueError(_ERR_FILTER_MUST_BE_RELATIVE.format(filter=self.filter))
+
+        return filter_normalized
+
+    def _resolve(self) -> tuple[str, str]:
+        """Resolve the source path and destination path.
+
+        Returns:
+            Tuple of (source_path, destination_relative_path)
+
+        Raises:
+            FileNotFoundError: If root or filter path does not exist.
+            ValueError: If filter is invalid (wildcards, absolute, escapes root, or applied to file).
+        """
+        filter_normalized = self._validate_filter()
+        root_normalized = os.path.normpath(os.path.abspath(self.root))
+
+        if filter_normalized is None:
+            if not os.path.exists(root_normalized):
+                raise FileNotFoundError(_ERR_ROOT_NOT_FOUND.format(root=self.root, resolved=root_normalized))
+            return root_normalized, os.path.basename(root_normalized)
+
+        if not os.path.exists(root_normalized):
+            raise FileNotFoundError(_ERR_ROOT_NOT_FOUND.format(root=self.root, resolved=root_normalized))
+
+        if os.path.isfile(root_normalized):
+            raise ValueError(_ERR_FILTER_ON_FILE_ROOT.format(root=self.root))
+
+        source = os.path.normpath(os.path.join(root_normalized, filter_normalized))
+
+        if not (source.startswith(root_normalized + os.sep) or source == root_normalized):
+            raise ValueError(_ERR_FILTER_ESCAPES_ROOT.format(filter=self.filter, root=self.root))
+
+        if not os.path.exists(source):
+            raise FileNotFoundError(_ERR_FILTER_NOT_FOUND.format(filter=self.filter, root=self.root, resolved=source))
+
+        return source, filter_normalized

snowflake/ml/model/custom_model.py

@@ -4,10 +4,13 @@ from typing import Any, Callable, Coroutine, Generator, Optional, Union
 
 import anyio
 import pandas as pd
-from typing_extensions import deprecated
+from typing_extensions import Concatenate, ParamSpec, deprecated
 
 from snowflake.ml.model import type_hints as model_types
 
+# Captures additional keyword-only parameters for inference methods
+InferenceParams = ParamSpec("InferenceParams")
+
 
 class MethodRef:
     """Represents a method invocation of an instance of `ModelRef`.
@@ -217,7 +220,7 @@ class CustomModel:
 
     def _get_infer_methods(
         self,
-    ) -> Generator[Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame], None, None]:
+    ) -> Generator[Callable[..., pd.DataFrame], None, None]:
         """Returns all methods in CLS with `inference_api` decorator as the outermost decorator."""
         for cls_method_str in dir(self):
             cls_method = getattr(self, cls_method_str)
@@ -240,7 +243,7 @@ class CustomModel:
         return rv
 
 
-def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]) -> None:
+def _validate_predict_function(func: Callable[..., pd.DataFrame]) -> None:
     """Validate the user provided predict method.
 
     Args:
@@ -248,19 +251,22 @@ def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.D
 
     Raises:
         TypeError: Raised when the method is not a callable object.
-        TypeError: Raised when the method does not have 2 arguments (self and X).
+        TypeError: Raised when the method does not have at least 2 arguments (self and X).
         TypeError: Raised when the method does not have typing annotation.
         TypeError: Raised when the method's input (X) does not have type pd.DataFrame.
         TypeError: Raised when the method's output does not have type pd.DataFrame.
+        TypeError: Raised when additional parameters are not keyword-only with defaults.
     """
     if not callable(func):
         raise TypeError("Predict method is not callable.")
 
     func_signature = inspect.signature(func)
-    if len(func_signature.parameters) != 2:
-        raise TypeError("Predict method should have exact 2 arguments.")
+    func_signature_params = list(func_signature.parameters.values())
+
+    if len(func_signature_params) < 2:
+        raise TypeError("Predict method should have at least 2 arguments.")
 
-    input_annotation = list(func_signature.parameters.values())[1].annotation
+    input_annotation = func_signature_params[1].annotation
     output_annotation = func_signature.return_annotation
 
     if input_annotation == inspect.Parameter.empty or output_annotation == inspect.Signature.empty:
@@ -275,17 +281,53 @@ def _validate_predict_function(func: Callable[[model_types.CustomModelType, pd.D
     ):
         raise TypeError("Output for predict method should have type pandas.DataFrame.")
 
+    # Validate additional parameters (beyond self and input) are keyword-only with defaults
+    for func_signature_param in func_signature_params[2:]:
+        _validate_parameter(func_signature_param)
+
+
+def _validate_parameter(param: inspect.Parameter) -> None:
+    """Validate a parameter."""
+    if param.kind != inspect.Parameter.KEYWORD_ONLY:
+        raise TypeError(f"Parameter '{param.name}' must be keyword-only (defined after '*' in signature).")
+    if param.default == inspect.Parameter.empty:
+        raise TypeError(f"Parameter '{param.name}' must have a default value.")
+    if param.annotation == inspect.Parameter.empty:
+        raise TypeError(f"Parameter '{param.name}' must have a type annotation.")
+
+    # Validate annotation is a supported type
+    supported_types = {int, float, str, bool}
+    if param.annotation not in supported_types:
+        raise TypeError(
+            f"Parameter '{param.name}' has unsupported type annotation '{param.annotation}'. "
+            f"Supported types are: int, float, str, bool"
+        )
+
+
+def get_method_parameters(func: Callable[..., Any]) -> list[tuple[str, Any, Any]]:
+    """Extract keyword-only parameters with defaults from an inference method.
+
+    Args:
+        func: The inference method.
+
+    Returns:
+        A list of tuples (name, type annotation, default value) for each keyword-only parameter.
+    """
+    func_signature = inspect.signature(func)
+    params = list(func_signature.parameters.values())
+    return [(param.name, param.annotation, param.default) for param in params[2:]]
+
 
 def inference_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
-) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func: Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame],
+) -> Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     return func
 
 
 def partitioned_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
-) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func: Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame],
+) -> Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     func.__dict__["_is_partitioned_api"] = True
     return func
@@ -296,8 +338,8 @@ def partitioned_api(
     " Use snowflake.ml.custom_model.partitioned_api instead."
 )
 def partitioned_inference_api(
-    func: Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame],
-) -> Callable[[model_types.CustomModelType, pd.DataFrame], pd.DataFrame]:
+    func: Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame],
+) -> Callable[Concatenate[model_types.CustomModelType, pd.DataFrame, InferenceParams], pd.DataFrame]:
     func.__dict__["_is_inference_api"] = True
     func.__dict__["_is_partitioned_api"] = True
     return func