snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

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

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

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

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

snowflake/ml/model/_packager/model_packager.py

@@ -1,6 +1,6 @@
 import os
 from types import ModuleType
-from typing import Dict, List, Optional
+from typing import Optional
 
 from absl import logging
 
@@ -38,15 +38,17 @@ class ModelPackager:
         *,
         name: str,
         model: model_types.SupportedModelType,
-        signatures: Optional[Dict[str, model_signature.ModelSignature]] = None,
+        signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
-        metadata: Optional[Dict[str, str]] = None,
-        conda_dependencies: Optional[List[str]] = None,
-        pip_requirements: Optional[List[str]] = None,
-        artifact_repository_map: Optional[Dict[str, str]] = None,
+        metadata: Optional[dict[str, str]] = None,
+        conda_dependencies: Optional[list[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
+        artifact_repository_map: Optional[dict[str, str]] = None,
+        resource_constraint: Optional[dict[str, str]] = None,
+        target_platforms: Optional[list[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
-        ext_modules: Optional[List[ModuleType]] = None,
-        code_paths: Optional[List[str]] = None,
+        ext_modules: Optional[list[ModuleType]] = None,
+        code_paths: Optional[list[str]] = None,
         options: model_types.ModelSaveOption,
         task: model_types.Task = model_types.Task.UNKNOWN,
     ) -> model_meta.ModelMetadata:
@@ -75,8 +77,10 @@ class ModelPackager:
             conda_dependencies=conda_dependencies,
             pip_requirements=pip_requirements,
             artifact_repository_map=artifact_repository_map,
+            resource_constraint=resource_constraint,
             python_version=python_version,
             task=task,
+            target_platforms=target_platforms,
             **options,
         ) as meta:
             model_blobs_path = os.path.join(self.local_dir_path, ModelPackager.MODEL_BLOBS_DIR)

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

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

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

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

snowflake/ml/model/_signatures/core.py

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

snowflake/ml/model/_signatures/dmatrix_handler.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, List, Literal, Optional, Sequence
+from typing import TYPE_CHECKING, Literal, Optional, Sequence
 
 import numpy as np
 import pandas as pd
@@ -45,7 +45,7 @@ class XGBoostDMatrixHandler(base_handler.BaseDataHandler["xgboost.DMatrix"]):
     @staticmethod
     def infer_signature(data: "xgboost.DMatrix", role: Literal["input", "output"]) -> Sequence[core.BaseFeatureSpec]:
         feature_prefix = f"{XGBoostDMatrixHandler.FEATURE_PREFIX}_"
-        features: List[core.BaseFeatureSpec] = []
+        features: list[core.BaseFeatureSpec] = []
         role_prefix = (
             XGBoostDMatrixHandler.INPUT_PREFIX if role == "input" else XGBoostDMatrixHandler.OUTPUT_PREFIX
         ) + "_"

snowflake/ml/model/_signatures/utils.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Any, Dict, List, Optional, Sequence
+from typing import Any, Optional, Sequence
 
 import numpy as np
 import numpy.typing as npt
@@ -12,7 +12,7 @@ from snowflake.ml._internal.exceptions import (
 from snowflake.ml.model._signatures import core
 
 
-def convert_list_to_ndarray(data: List[Any]) -> npt.NDArray[Any]:
+def convert_list_to_ndarray(data: list[Any]) -> npt.NDArray[Any]:
     """Create a numpy array from list or nested list. Avoid ragged list and unaligned types.
 
     Args:
@@ -49,7 +49,7 @@ def convert_list_to_ndarray(data: List[Any]) -> npt.NDArray[Any]:
 
 
 def rename_features(
-    features: Sequence[core.BaseFeatureSpec], feature_names: Optional[List[str]] = None
+    features: Sequence[core.BaseFeatureSpec], feature_names: Optional[list[str]] = None
 ) -> Sequence[core.BaseFeatureSpec]:
     """It renames the feature in features provided optional feature names.
 
@@ -104,7 +104,7 @@ def rename_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureSpec
     return data
 
 
-def huggingface_pipeline_signature_auto_infer(task: str, params: Dict[str, Any]) -> Optional[core.ModelSignature]:
+def huggingface_pipeline_signature_auto_infer(task: str, params: dict[str, Any]) -> Optional[core.ModelSignature]:
     # Text
 
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.ConversationalPipeline
@@ -351,7 +351,7 @@ def series_dropna(series: pd.Series) -> pd.Series:
     return series.dropna(inplace=False).reset_index(drop=True).convert_dtypes()
 
 
-def infer_list(name: str, data: List[Any]) -> core.BaseFeatureSpec:
+def infer_list(name: str, data: list[Any]) -> core.BaseFeatureSpec:
     """Infer the feature specification from a list.
 
     Args:
@@ -382,7 +382,7 @@ def infer_list(name: str, data: List[Any]) -> core.BaseFeatureSpec:
     return core.FeatureSpec(name=name, dtype=arr_dtype, shape=arr.shape)
 
 
-def infer_dict(name: str, data: Dict[str, Any]) -> core.FeatureGroupSpec:
+def infer_dict(name: str, data: dict[str, Any]) -> core.FeatureGroupSpec:
     """Infer the feature specification from a dictionary.
 
     Args:

snowflake/ml/model/custom_model.py

@@ -1,6 +1,6 @@
 import functools
 import inspect
-from typing import Any, Callable, Coroutine, Dict, Generator, List, Optional, Union
+from typing import Any, Callable, Coroutine, Generator, Optional, Union
 
 import anyio
 import pandas as pd
@@ -78,7 +78,7 @@ class ModelRef:
             return MethodRef(self, method_name)
         raise AttributeError(f"Method {method_name} not found in model {self._name}.")
 
-    def __getstate__(self) -> Dict[str, Any]:
+    def __getstate__(self) -> dict[str, Any]:
         state = self.__dict__.copy()
         del state["_model"]
         return state
@@ -113,8 +113,8 @@ class ModelContext:
     def __init__(
         self,
         *,
-        artifacts: Optional[Union[Dict[str, str], str, model_types.SupportedModelType]] = None,
-        models: Optional[Union[Dict[str, model_types.SupportedModelType], str, model_types.SupportedModelType]] = None,
+        artifacts: Optional[Union[dict[str, str], str, model_types.SupportedModelType]] = None,
+        models: Optional[Union[dict[str, model_types.SupportedModelType], str, model_types.SupportedModelType]] = None,
         **kwargs: Optional[Union[str, model_types.SupportedModelType]],
     ) -> None:
         """Initialize the model context.
@@ -130,8 +130,8 @@ class ModelContext:
             ValueError: Raised when the model name is duplicated.
         """
 
-        self.artifacts: Dict[str, str] = dict()
-        self.model_refs: Dict[str, ModelRef] = dict()
+        self.artifacts: dict[str, str] = dict()
+        self.model_refs: dict[str, ModelRef] = dict()
 
         # In case that artifacts is a dictionary, assume the original usage,
         # which is to pass in a dictionary of artifacts.
@@ -185,7 +185,7 @@ class ModelContext:
         return self.model_refs[name]
 
     def __getitem__(self, key: str) -> Union[str, ModelRef]:
-        combined: Dict[str, Union[str, ModelRef]] = {**self.artifacts, **self.model_refs}
+        combined: dict[str, Union[str, ModelRef]] = {**self.artifacts, **self.model_refs}
         if key not in combined:
             raise KeyError(f"Key {key} not found in the kwargs, current available keys are: {combined.keys()}")
         return combined[key]
@@ -226,7 +226,7 @@ class CustomModel:
                 else:
                     raise TypeError("A non-method inference API function is not supported.")
 
-    def _get_partitioned_infer_methods(self) -> List[str]:
+    def _get_partitioned_infer_methods(self) -> list[str]:
         """Returns all methods in CLS with `partitioned_inference_api` as the outermost decorator."""
         rv = []
         for cls_method_str in dir(self):

snowflake/ml/model/model_signature.py

@@ -1,18 +1,7 @@
 import enum
 import json
 import warnings
-from typing import (
-    Any,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Sequence,
-    Tuple,
-    Type,
-    Union,
-    cast,
-)
+from typing import Any, Literal, Optional, Sequence, Union, cast
 
 import numpy as np
 import pandas as pd
@@ -30,7 +19,7 @@ from snowflake.ml._internal.utils import formatting, identifier, sql_identifier
 from snowflake.ml.model import type_hints as model_types
 from snowflake.ml.model._signatures import (
     base_handler,
-    builtins_handler as builtins_handler,
+    builtins_handler,
     core,
     dmatrix_handler,
     numpy_handler,
@@ -48,7 +37,7 @@ FeatureGroupSpec = core.FeatureGroupSpec
 ModelSignature = core.ModelSignature
 
 
-_LOCAL_DATA_HANDLERS: List[Type[base_handler.BaseDataHandler[Any]]] = [
+_LOCAL_DATA_HANDLERS: list[type[base_handler.BaseDataHandler[Any]]] = [
     pandas_handler.PandasDataFrameHandler,
     numpy_handler.NumpyArrayHandler,
     builtins_handler.ListOfBuiltinHandler,
@@ -414,7 +403,7 @@ class SnowparkIdentifierRule(enum.Enum):
 
 def _get_dataframe_values_range(
     df: snowflake.snowpark.DataFrame,
-) -> Dict[str, Union[Tuple[int, int], Tuple[float, float]]]:
+) -> dict[str, Union[tuple[int, int], tuple[float, float]]]:
     columns = [
         F.array_construct(F.min(field.name), F.max(field.name)).as_(field.name)
         for field in df.schema.fields
@@ -429,7 +418,7 @@ def _get_dataframe_values_range(
             original_exception=ValueError(f"Unable to get the value range of fields {df.columns}"),
         )
     return cast(
-        Dict[str, Union[Tuple[int, int], Tuple[float, float]]],
+        dict[str, Union[tuple[int, int], tuple[float, float]]],
         {
             sql_identifier.SqlIdentifier(k, case_sensitive=True).identifier(): (json.loads(v)[0], json.loads(v)[1])
             for k, v in res[0].as_dict().items()
@@ -456,7 +445,7 @@ def _validate_snowpark_data(
         - inferred: signature `a` - Snowpark DF `"a"`, use `get_inferred_name`
         - normalized: signature `a` - Snowpark DF `A`, use `resolve_identifier`
     """
-    errors: Dict[SnowparkIdentifierRule, List[Exception]] = {
+    errors: dict[SnowparkIdentifierRule, list[Exception]] = {
         SnowparkIdentifierRule.INFERRED: [],
         SnowparkIdentifierRule.NORMALIZED: [],
     }
@@ -549,7 +538,7 @@ def _validate_snowpark_type_feature(
     field: spt.StructField,
     ft_type: DataType,
     ft_name: str,
-    value_range: Optional[Union[Tuple[int, int], Tuple[float, float]]],
+    value_range: Optional[Union[tuple[int, int], tuple[float, float]]],
     strict: bool = False,
 ) -> None:
     field_data_type = field.datatype
@@ -716,8 +705,8 @@ def _convert_and_validate_local_data(
 def infer_signature(
     input_data: model_types.SupportedLocalDataType,
     output_data: model_types.SupportedLocalDataType,
-    input_feature_names: Optional[List[str]] = None,
-    output_feature_names: Optional[List[str]] = None,
+    input_feature_names: Optional[list[str]] = None,
+    output_feature_names: Optional[list[str]] = None,
     input_data_limit: Optional[int] = 100,
     output_data_limit: Optional[int] = 100,
 ) -> core.ModelSignature:

snowflake/ml/model/models/huggingface_pipeline.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from packaging import version
 
@@ -13,7 +13,7 @@ class HuggingFacePipelineModel:
         revision: Optional[str] = None,
         token: Optional[str] = None,
         trust_remote_code: Optional[bool] = None,
-        model_kwargs: Optional[Dict[str, Any]] = None,
+        model_kwargs: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -65,6 +65,7 @@ class HuggingFacePipelineModel:
             warnings.warn(
                 "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.",
                 FutureWarning,
+                stacklevel=2,
             )
             if token is not None:
                 raise ValueError(
@@ -183,7 +184,8 @@ class HuggingFacePipelineModel:
             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."
+                "Using a pipeline without specifying a model name and revision in production is not recommended.",
+                stacklevel=2,
             )
             if config is None and isinstance(model, str):
                 config_obj = transformers.AutoConfig.from_pretrained(
@@ -200,7 +202,8 @@ class HuggingFacePipelineModel:
             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`."
+                    " will most likely encounter unexpected behavior. Please remove `device` and keep `device_map`.",
+                    stacklevel=2,
                 )
 
         # ==== End pipeline logic from transformers ====

snowflake/ml/model/type_hints.py

@@ -1,6 +1,6 @@
 # mypy: disable-error-code="import"
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, Literal, Sequence, TypedDict, TypeVar, Union
+from typing import TYPE_CHECKING, Literal, Sequence, TypedDict, TypeVar, Union
 
 import numpy.typing as npt
 from typing_extensions import NotRequired
@@ -32,7 +32,7 @@ _SupportedBuiltins = Union[
     bool,
     str,
     bytes,
-    Dict[str, Union["_SupportedBuiltins", "_SupportedBuiltinsList"]],
+    dict[str, Union["_SupportedBuiltins", "_SupportedBuiltinsList"]],
     "_SupportedBuiltinsList",
 ]
 _SupportedNumpyDtype = Union[
@@ -147,13 +147,15 @@ class BaseModelSaveOption(TypedDict):
     embed_local_ml_library: Embedding local SnowML into the code directory of the folder.
     relax_version: Whether or not relax the version constraints of the dependencies if unresolvable in Warehouse.
         It detects any ==x.y.z in specifiers and replaced with >=x.y, <(x+1). Defaults to True.
+    save_location: Local directory path to save the model and metadata.
     """
 
     embed_local_ml_library: NotRequired[bool]
     relax_version: NotRequired[bool]
     function_type: NotRequired[Literal["FUNCTION", "TABLE_FUNCTION"]]
-    method_options: NotRequired[Dict[str, ModelMethodSaveOptions]]
+    method_options: NotRequired[dict[str, ModelMethodSaveOptions]]
     enable_explainability: NotRequired[bool]
+    save_location: NotRequired[str]
 
 
 class CatBoostModelSaveOptions(BaseModelSaveOption):

snowflake/ml/modeling/_internal/estimator_utils.py

@@ -1,7 +1,7 @@
 import inspect
 import numbers
 import os
-from typing import Any, Callable, Dict, List, Set, Tuple
+from typing import Any, Callable
 
 import cloudpickle as cp
 import numpy as np
@@ -16,7 +16,7 @@ from snowflake.snowpark import Session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
-def validate_sklearn_args(args: Dict[str, Tuple[Any, Any, bool]], klass: type) -> Dict[str, Any]:
+def validate_sklearn_args(args: dict[str, tuple[Any, Any, bool]], klass: type) -> dict[str, Any]:
     """Validate if all the keyword args are supported by current version of SKLearn/XGBoost object.
 
     Args:
@@ -71,7 +71,7 @@ def transform_snowml_obj_to_sklearn_obj(obj: Any) -> Any:
         return obj
 
 
-def gather_dependencies(obj: Any) -> Set[str]:
+def gather_dependencies(obj: Any) -> set[str]:
     """Gathers dependencies from the SnowML Estimator and Transformer objects.
 
     Args:
@@ -82,7 +82,7 @@ def gather_dependencies(obj: Any) -> Set[str]:
     """
 
     if isinstance(obj, list) or isinstance(obj, tuple):
-        deps: Set[str] = set()
+        deps: set[str] = set()
         for elem in obj:
             deps = deps | set(gather_dependencies(elem))
         return deps
@@ -167,8 +167,8 @@ def get_module_name(model: object) -> str:
 
 
 def handle_inference_result(
-    inference_res: Any, output_cols: List[str], inference_method: str, within_udf: bool = False
-) -> Tuple[npt.NDArray[Any], List[str]]:
+    inference_res: Any, output_cols: list[str], inference_method: str, within_udf: bool = False
+) -> tuple[npt.NDArray[Any], list[str]]:
     if isinstance(inference_res, list) and len(inference_res) > 0 and isinstance(inference_res[0], np.ndarray):
         # In case of multioutput estimators, predict_proba, decision_function etc., functions return a list of
         # ndarrays. We need to concatenate them.
@@ -248,7 +248,7 @@ def create_temp_stage(session: Session) -> str:
 
 
 def upload_model_to_stage(
-    stage_name: str, estimator: object, session: Session, statement_params: Dict[str, str]
+    stage_name: str, estimator: object, session: Session, statement_params: dict[str, str]
 ) -> str:
     """Util method to pickle and upload the model to a temp Snowflake stage.