snowflake-ml-python 1.20.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

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/compute_pool.py

@@ -0,0 +1,2 @@
+DEFAULT_CPU_COMPUTE_POOL = "SYSTEM_COMPUTE_POOL_CPU"
+DEFAULT_GPU_COMPUTE_POOL = "SYSTEM_COMPUTE_POOL_GPU"

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

snowflake/ml/model/model_signature.py

@@ -34,6 +34,9 @@ DataType = core.DataType
 BaseFeatureSpec = core.BaseFeatureSpec
 FeatureSpec = core.FeatureSpec
 FeatureGroupSpec = core.FeatureGroupSpec
+BaseParamSpec = core.BaseParamSpec
+ParamSpec = core.ParamSpec
+ParamGroupSpec = core.ParamGroupSpec
 ModelSignature = core.ModelSignature
 
 
@@ -711,6 +714,7 @@ def infer_signature(
     output_feature_names: Optional[list[str]] = None,
     input_data_limit: Optional[int] = 100,
     output_data_limit: Optional[int] = 100,
+    params: Optional[Sequence[core.BaseParamSpec]] = None,
 ) -> core.ModelSignature:
     """
     Infer model signature from given input and output sample data.
@@ -740,12 +744,20 @@ def infer_signature(
         output_data_limit: Limit the number of rows to be used in signature inference in the output data. Defaults to
             100. If None, all rows are used. If the number of rows in the output data is less than the limit, all rows
             are used.
+        params: Optional sequence of parameter specifications to include in the signature. Parameters define
+            optional configuration values that can be passed to model inference. Defaults to None.
+
+    Raises:
+        SnowflakeMLException: ValueError: Raised when input data contains columns matching parameter names.
 
     Returns:
         A model signature inferred from the given input and output sample data.
+
+    # noqa: DAR402
     """
     inputs = _infer_signature(_truncate_data(input_data, input_data_limit), role="input")
     inputs = utils.rename_features(inputs, input_feature_names)
+
     outputs = _infer_signature(_truncate_data(output_data, output_data_limit), role="output")
     outputs = utils.rename_features(outputs, output_feature_names)
-    return core.ModelSignature(inputs, outputs)
+    return core.ModelSignature(inputs, outputs, params=params)

snowflake/ml/model/models/huggingface.py

@@ -0,0 +1,285 @@
+import logging
+import warnings
+from typing import Any, Optional, Union
+
+from packaging import version
+
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model.compute_pool import DEFAULT_CPU_COMPUTE_POOL
+
+logger = logging.getLogger(__name__)
+
+
+_TELEMETRY_PROJECT = "MLOps"
+_TELEMETRY_SUBPROJECT = "ModelManagement"
+
+
+class TransformersPipeline:
+    def __init__(
+        self,
+        task: Optional[str] = None,
+        model: Optional[str] = None,
+        *,
+        revision: Optional[str] = None,
+        token_or_secret: Optional[str] = None,
+        trust_remote_code: Optional[bool] = None,
+        model_kwargs: Optional[dict[str, Any]] = None,
+        compute_pool_for_log: Optional[str] = DEFAULT_CPU_COMPUTE_POOL,
+        # repo snapshot download args
+        allow_patterns: Optional[Union[list[str], str]] = None,
+        ignore_patterns: Optional[Union[list[str], str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Utility factory method to build a wrapper over transformers [`Pipeline`].
+        When deploying, this wrapper will create a real pipeline object and loading tokenizers and models.
+
+        For pipelines docs, please refer:
+        https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.pipeline
+
+        Args:
+            task: The task that pipeline will be used. If None it would be inferred from model.
+                For available tasks, please refer Transformers's documentation. Defaults to None.
+            model: The model that will be used by the pipeline to make predictions. This can only be a model identifier
+                currently. If not provided, the default for the `task` will be loaded. Defaults to None.
+            revision: When passing a task name or a string model identifier: The specific model version to use. It can
+                be a branch name, a tag name, or a commit id, since we use a git-based system for storing models and
+                other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. Defaults to None.
+            token_or_secret: The token to use as HTTP bearer authorization for remote files. Defaults to None.
+              The token can be a token or a secret. If a secret is provided, it must a fully qualified secret name.
+            trust_remote_code: Whether or not to allow for custom code defined on the Hub in their own modeling,
+                configuration, tokenization or even pipeline files. This option should only be set to `True` for
+                repositories you trust and in which you have read the code, as it will execute code present on the Hub.
+                Defaults to None.
+            model_kwargs: Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,`.
+                Defaults to None.
+            compute_pool_for_log: The compute pool to use for logging the model. Defaults to DEFAULT_CPU_COMPUTE_POOL.
+                If a string is provided, it will be used as the compute pool name. This override allows for logging
+                the model when there is no system compute pool available.
+                If None is passed,
+                    if `huggingface_hub` is installed, the model artifacts will be downloaded
+                    from the HuggingFace repository.
+                    otherwise, the only the metadata will be logged to snowflake.
+            allow_patterns: If provided, only files matching at least one pattern are downloaded.
+            ignore_patterns: If provided, files matching any of the patterns are not downloaded.
+            kwargs: Additional keyword arguments passed along to the specific pipeline init (see the documentation for
+                the corresponding pipeline class for possible values).
+
+        Raises:
+            RuntimeError: Raised when the input argument cannot determine the pipeline.
+            ValueError: Raised when the pipeline contains remote code but trust_remote_code is not set or False.
+            ValueError: Raised when having conflicting arguments.
+
+        .. # noqa: DAR003
+        """
+        import transformers
+
+        config = kwargs.get("config", None)
+        tokenizer = kwargs.get("tokenizer", None)
+        framework = kwargs.get("framework", None)
+        feature_extractor = kwargs.get("feature_extractor", None)
+
+        self.secret_identifier: Optional[str] = None
+        uses_secret = False
+        if token_or_secret is not None and isinstance(token_or_secret, str):
+            db, schema, secret_name = sql_identifier.parse_fully_qualified_name(token_or_secret)
+            if db is not None and schema is not None and secret_name is not None:
+                self.secret_identifier = sql_identifier.get_fully_qualified_name(
+                    db=db,
+                    schema=schema,
+                    object=secret_name,
+                )
+                uses_secret = True
+            else:
+                logger.info("The token_or_secret is not a fully qualified secret name. It will be used as is.")
+
+        can_download_snapshot = False
+        if compute_pool_for_log is None:
+            try:
+                import huggingface_hub as hf_hub
+
+                can_download_snapshot = True
+            except ImportError:
+                pass
+
+        if compute_pool_for_log is None and not can_download_snapshot:
+            logger.info(
+                "The model will be logged with metadata only. No model artifacts will be downloaded. "
+                "During deployment, the model artifacts will be downloaded from the HuggingFace repository."
+            )
+
+        # ==== Start pipeline logic from transformers ====
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        use_auth_token = model_kwargs.pop("use_auth_token", None)
+        if use_auth_token is not None:
+            warnings.warn(
+                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            if token_or_secret is not None:
+                raise ValueError(
+                    "`token_or_secret` and `use_auth_token` are both specified. "
+                    "Please set only the argument `token_or_secret`."
+                )
+            token_or_secret = use_auth_token
+
+        hub_kwargs = {
+            "revision": revision,
+            "token": token_or_secret,
+            "trust_remote_code": trust_remote_code,
+            "_commit_hash": None,
+        }
+
+        # Backward compatibility since HF interface change.
+        if version.parse(transformers.__version__) < version.parse("4.32.0"):
+            # Backward compatibility since HF interface change.
+            hub_kwargs["use_auth_token"] = hub_kwargs["token"]
+            del hub_kwargs["token"]
+
+        if task is None and model is None:
+            raise RuntimeError(
+                "Impossible to instantiate a pipeline without either a task or a model being specified. "
+            )
+
+        if model is None and tokenizer is not None:
+            raise RuntimeError(
+                "Impossible to instantiate a pipeline with tokenizer specified but not the model as the provided"
+                " tokenizer may not be compatible with the default model. Please provide an identifier to a pretrained"
+                " model when providing tokenizer."
+            )
+        if model is None and feature_extractor is not None:
+            raise RuntimeError(
+                "Impossible to instantiate a pipeline with feature_extractor specified but not the model as the "
+                "provided feature_extractor may not be compatible with the default model. Please provide an identifier"
+                " to a pretrained model when providing feature_extractor."
+            )
+
+        # ==== End pipeline logic from transformers ====
+
+        # We only support string as model argument.
+
+        if model is not None and not isinstance(model, str):
+            raise RuntimeError(f"Impossible to use non-string model as input for class {self.__class__.__name__}.")
+
+        # ==== Start pipeline logic (Config) from transformers ====
+
+        # Config is the primordial information item.
+        # Instantiate config if needed
+        config_obj = None
+
+        if not can_download_snapshot:
+            if isinstance(config, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    config, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+            elif config is None and isinstance(model, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+            # We only support string as config argument.
+            elif config is not None and not isinstance(config, str):
+                raise RuntimeError(
+                    f"Impossible to use non-string config as input for class {self.__class__.__name__}. "
+                    "Use transformers.Pipeline object if required."
+                )
+
+        # ==== Start pipeline logic (Task) from transformers ====
+
+        custom_tasks = {}
+        if config_obj is not None and len(getattr(config_obj, "custom_pipelines", {})) > 0:
+            custom_tasks = config_obj.custom_pipelines
+            if task is None and trust_remote_code is not False:
+                if len(custom_tasks) == 1:
+                    task = list(custom_tasks.keys())[0]
+                else:
+                    raise RuntimeError(
+                        "We can't infer the task automatically for this model as there are multiple tasks available. "
+                        f"Pick one in {', '.join(custom_tasks.keys())}"
+                    )
+
+        if task is None and model is not None:
+            task = transformers.pipelines.get_task(model, token_or_secret)
+
+        # Retrieve the task
+        if task in custom_tasks:
+            normalized_task = task
+            targeted_task, task_options = transformers.pipelines.clean_custom_task(custom_tasks[task])
+            if not trust_remote_code:
+                raise ValueError(
+                    "Loading this pipeline requires you to execute the code in the pipeline file in that"
+                    " repo on your local machine. Make sure you have read the code there to avoid malicious use, then"
+                    " set the option `trust_remote_code=True` to remove this error."
+                )
+        else:
+            (
+                normalized_task,
+                targeted_task,
+                task_options,
+            ) = transformers.pipelines.check_task(task)
+
+        # ==== Start pipeline logic (Model) from transformers ====
+
+        # Use default model/config/tokenizer for the task if no model is provided
+        if model is None:
+            # At that point framework might still be undetermined
+            (
+                model,
+                default_revision,
+            ) = transformers.pipelines.get_default_model_and_revision(targeted_task, framework, task_options)
+            revision = revision if revision is not None else default_revision
+            warnings.warn(
+                f"No model was supplied, defaulted to {model} and revision"
+                f" {revision} ({transformers.pipelines.HUGGINGFACE_CO_RESOLVE_ENDPOINT}/{model}).\n"
+                "Using a pipeline without specifying a model name and revision in production is not recommended.",
+                stacklevel=2,
+            )
+            if not can_download_snapshot and config is None and isinstance(model, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+
+        if kwargs.get("device_map", None) is not None:
+            if "device_map" in model_kwargs:
+                raise ValueError(
+                    'You cannot use both `pipeline(... device_map=..., model_kwargs={"device_map":...})` as those'
+                    " arguments might conflict, use only one.)"
+                )
+            if kwargs.get("device", None) is not None:
+                warnings.warn(
+                    "Both `device` and `device_map` are specified. `device` will override `device_map`. You"
+                    " will most likely encounter unexpected behavior. Please remove `device` and keep `device_map`.",
+                    stacklevel=2,
+                )
+
+        repo_snapshot_dir: Optional[str] = None
+        if can_download_snapshot and not uses_secret:
+            try:
+                repo_snapshot_dir = hf_hub.snapshot_download(
+                    repo_id=model,
+                    revision=revision,
+                    token=token_or_secret,
+                    allow_patterns=allow_patterns,
+                    ignore_patterns=ignore_patterns,
+                )
+            except ImportError:
+                logger.info("huggingface_hub package is not installed, skipping snapshot download")
+
+        # ==== End pipeline logic from transformers ====
+
+        self.model = model
+        self.task = normalized_task
+        self.revision = revision
+        self.token_or_secret = token_or_secret
+        self.trust_remote_code = trust_remote_code
+        self.model_kwargs = model_kwargs
+        self.tokenizer = tokenizer
+
+        self.repo_snapshot_dir = repo_snapshot_dir
+        self.compute_pool_for_log = compute_pool_for_log
+        self.__dict__.update(kwargs)