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

snowflake/ml/modeling/_internal/snowpark_implementations/xgboost_external_memory_trainer.py

@@ -1,7 +1,7 @@
 import inspect
 import os
 import tempfile
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 import cloudpickle as cp
 import pandas as pd
@@ -41,13 +41,13 @@ _PROJECT = "ModelDevelopment"
 
 
 def get_data_iterator(
-    file_paths: List[str],
+    file_paths: list[str],
     batch_size: int,
-    input_cols: List[str],
-    label_cols: List[str],
+    input_cols: list[str],
+    label_cols: list[str],
     sample_weight_col: Optional[str] = None,
 ) -> Any:
-    from typing import List, Optional
+    from typing import Optional
 
     import xgboost
 
@@ -60,10 +60,10 @@ def get_data_iterator(
 
         def __init__(
             self,
-            file_paths: List[str],
+            file_paths: list[str],
             batch_size: int,
-            input_cols: List[str],
-            label_cols: List[str],
+            input_cols: list[str],
+            label_cols: list[str],
             sample_weight_col: Optional[str] = None,
         ) -> None:
             """
@@ -151,10 +151,10 @@ def get_data_iterator(
 
 def train_xgboost_model(
     estimator: object,
-    file_paths: List[str],
+    file_paths: list[str],
     batch_size: int,
-    input_cols: List[str],
-    label_cols: List[str],
+    input_cols: list[str],
+    label_cols: list[str],
     sample_weight_col: Optional[str] = None,
 ) -> object:
     """
@@ -247,8 +247,8 @@ class XGBoostExternalMemoryTrainer(SnowparkModelTrainer):
         estimator: object,
         dataset: DataFrame,
         session: Session,
-        input_cols: List[str],
-        label_cols: Optional[List[str]],
+        input_cols: list[str],
+        label_cols: Optional[list[str]],
         sample_weight_col: Optional[str],
         autogenerated: bool = False,
         subproject: str = "",
@@ -285,8 +285,8 @@ class XGBoostExternalMemoryTrainer(SnowparkModelTrainer):
         self,
         model_spec: ModelSpecifications,
         session: Session,
-        statement_params: Dict[str, str],
-        import_file_paths: List[str],
+        statement_params: dict[str, str],
+        import_file_paths: list[str],
     ) -> Any:
         fit_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
 
@@ -308,10 +308,10 @@ class XGBoostExternalMemoryTrainer(SnowparkModelTrainer):
             session: Session,
             dataset_stage_name: str,
             batch_size: int,
-            input_cols: List[str],
-            label_cols: List[str],
+            input_cols: list[str],
+            label_cols: list[str],
             sample_weight_col: Optional[str],
-            statement_params: Dict[str, str],
+            statement_params: dict[str, str],
         ) -> str:
             import os
             import sys
@@ -365,7 +365,7 @@ class XGBoostExternalMemoryTrainer(SnowparkModelTrainer):
 
         return fit_wrapper_sproc
 
-    def _write_training_data_to_stage(self, dataset_stage_name: str) -> List[str]:
+    def _write_training_data_to_stage(self, dataset_stage_name: str) -> list[str]:
         """
         Materializes the training to the specified stage and returns the list of stage file paths.
 

snowflake/ml/modeling/_internal/transformer_protocols.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Optional, Protocol, TypedDict, Union
+from typing import Any, Optional, Protocol, TypedDict, Union
 
 import pandas as pd
 
@@ -29,9 +29,9 @@ class LocalModelTransformHandlers(Protocol):
     def batch_inference(
         self,
         inference_method: str,
-        input_cols: List[str],
-        expected_output_cols: List[str],
-        snowpark_input_cols: Optional[List[str]],
+        input_cols: list[str],
+        expected_output_cols: list[str],
+        snowpark_input_cols: Optional[list[str]],
         drop_input_cols: Optional[bool] = False,
         *args: Any,
         **kwargs: Any,
@@ -57,8 +57,8 @@ class LocalModelTransformHandlers(Protocol):
 
     def score(
         self,
-        input_cols: List[str],
-        label_cols: List[str],
+        input_cols: list[str],
+        label_cols: list[str],
         sample_weight_col: Optional[str],
         *args: Any,
         **kwargs: Any,
@@ -105,10 +105,10 @@ class RemoteModelTransformHandlers(Protocol):
     def batch_inference(
         self,
         inference_method: str,
-        input_cols: List[str],
-        expected_output_cols: List[str],
+        input_cols: list[str],
+        expected_output_cols: list[str],
         session: snowpark.Session,
-        dependencies: List[str],
+        dependencies: list[str],
         drop_input_cols: Optional[bool] = False,
         expected_output_cols_type: Optional[str] = "",
         *args: Any,
@@ -137,11 +137,11 @@ class RemoteModelTransformHandlers(Protocol):
 
     def score(
         self,
-        input_cols: List[str],
-        label_cols: List[str],
+        input_cols: list[str],
+        label_cols: list[str],
         session: snowpark.Session,
-        dependencies: List[str],
-        score_sproc_imports: List[str],
+        dependencies: list[str],
+        score_sproc_imports: list[str],
         sample_weight_col: Optional[str] = None,
         *args: Any,
         **kwargs: Any,
@@ -173,10 +173,10 @@ ModelTransformHandlers = Union[LocalModelTransformHandlers, RemoteModelTransform
 class BatchInferenceKwargsTypedDict(TypedDict, total=False):
     """A typed dict specifying all possible optional keyword args accepted by batch_inference() methods."""
 
-    snowpark_input_cols: Optional[List[str]]
+    snowpark_input_cols: Optional[list[str]]
     drop_input_cols: Optional[bool]
     session: snowpark.Session
-    dependencies: List[str]
+    dependencies: list[str]
     expected_output_cols_type: str
     n_neighbors: Optional[int]
     return_distance: bool
@@ -186,5 +186,5 @@ class ScoreKwargsTypedDict(TypedDict, total=False):
     """A typed dict specifying all possible optional keyword args accepted by score() methods."""
 
     session: snowpark.Session
-    dependencies: List[str]
-    score_sproc_imports: List[str]
+    dependencies: list[str]
+    score_sproc_imports: list[str]

snowflake/ml/modeling/framework/_utils.py

@@ -3,7 +3,7 @@
 import inspect
 import warnings
 from enum import Enum
-from typing import Any, Callable, Dict, Iterable, Optional, Union
+from typing import Any, Callable, Iterable, Optional, Union
 
 import numpy as np
 import sklearn
@@ -62,7 +62,7 @@ class BasicStatistics(str, Enum):
     MODE = "mode"
 
 
-def get_default_args(func: Callable[..., None]) -> Dict[str, Any]:
+def get_default_args(func: Callable[..., None]) -> dict[str, Any]:
     signature = inspect.signature(func)
     return {k: v.default for k, v in signature.parameters.items() if v.default is not inspect.Parameter.empty}
 
@@ -72,16 +72,16 @@ def generate_value_with_prefix(prefix: str) -> str:
 
 
 def get_filtered_valid_sklearn_args(
-    args: Dict[str, Any],
-    default_sklearn_args: Dict[str, Any],
+    args: dict[str, Any],
+    default_sklearn_args: dict[str, Any],
     sklearn_initial_keywords: Optional[Union[str, Iterable[str]]] = None,
     sklearn_unused_keywords: Optional[Union[str, Iterable[str]]] = None,
     snowml_only_keywords: Optional[Union[str, Iterable[str]]] = None,
-    sklearn_added_keyword_to_version_dict: Optional[Dict[str, str]] = None,
-    sklearn_added_kwarg_value_to_version_dict: Optional[Dict[str, Dict[str, str]]] = None,
-    sklearn_deprecated_keyword_to_version_dict: Optional[Dict[str, str]] = None,
-    sklearn_removed_keyword_to_version_dict: Optional[Dict[str, str]] = None,
-) -> Dict[str, Any]:
+    sklearn_added_keyword_to_version_dict: Optional[dict[str, str]] = None,
+    sklearn_added_kwarg_value_to_version_dict: Optional[dict[str, dict[str, str]]] = None,
+    sklearn_deprecated_keyword_to_version_dict: Optional[dict[str, str]] = None,
+    sklearn_removed_keyword_to_version_dict: Optional[dict[str, str]] = None,
+) -> dict[str, Any]:
     """
     Get valid sklearn keyword arguments with non-default values.
 
@@ -241,7 +241,7 @@ def to_native_format(obj: Any) -> Any:
     return obj.to_sklearn()
 
 
-def table_exists(session: snowpark.Session, table_name: str, statement_params: Dict[str, Any]) -> bool:
+def table_exists(session: snowpark.Session, table_name: str, statement_params: dict[str, Any]) -> bool:
     try:
         session.table(table_name).limit(0).collect(statement_params=statement_params)
         return True

snowflake/ml/modeling/framework/base.py

@@ -2,7 +2,7 @@
 import inspect
 from abc import abstractmethod
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Union, overload
+from typing import Any, Iterable, Mapping, Optional, Union, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -28,9 +28,9 @@ SKLEARN_SUPERVISED_ESTIMATORS = ["regressor", "classifier"]
 SKLEARN_SINGLE_OUTPUT_ESTIMATORS = ["DensityEstimator", "clusterer", "outlier_detector"]
 
 
-def _process_cols(cols: Optional[Union[str, Iterable[str]]]) -> List[str]:
+def _process_cols(cols: Optional[Union[str, Iterable[str]]]) -> list[str]:
     """Convert cols to a list."""
-    col_list: List[str] = []
+    col_list: list[str] = []
     if cols is None:
         return col_list
     elif type(cols) is list:
@@ -55,10 +55,10 @@ class Base:
             passthrough_cols: List columns not to be used or modified by the estimator/transformers.
                 These columns will be passed through all the estimator/transformer operations without any modifications.
         """
-        self.input_cols: List[str] = []
-        self.output_cols: List[str] = []
-        self.label_cols: List[str] = []
-        self.passthrough_cols: List[str] = []
+        self.input_cols: list[str] = []
+        self.output_cols: list[str] = []
+        self.label_cols: list[str] = []
+        self.passthrough_cols: list[str] = []
 
     def _create_unfitted_sklearn_object(self) -> Any:
         raise NotImplementedError()
@@ -66,7 +66,7 @@ class Base:
     def _create_sklearn_object(self) -> Any:
         raise NotImplementedError()
 
-    def get_input_cols(self) -> List[str]:
+    def get_input_cols(self) -> list[str]:
         """
         Input columns getter.
 
@@ -88,7 +88,7 @@ class Base:
         self.input_cols = _process_cols(input_cols)
         return self
 
-    def get_output_cols(self) -> List[str]:
+    def get_output_cols(self) -> list[str]:
         """
         Output columns getter.
 
@@ -110,7 +110,7 @@ class Base:
         self.output_cols = _process_cols(output_cols)
         return self
 
-    def get_label_cols(self) -> List[str]:
+    def get_label_cols(self) -> list[str]:
         """
         Label column getter.
 
@@ -132,7 +132,7 @@ class Base:
         self.label_cols = _process_cols(label_cols)
         return self
 
-    def get_passthrough_cols(self) -> List[str]:
+    def get_passthrough_cols(self) -> list[str]:
         """
         Passthrough columns getter.
 
@@ -215,7 +215,7 @@ class Base:
             )
 
     @classmethod
-    def _get_param_names(cls) -> List[str]:
+    def _get_param_names(cls) -> list[str]:
         """Get parameter names for the transformer"""
         # fetch the constructor or the original constructor before
         # deprecation wrapping if any
@@ -244,7 +244,7 @@ class Base:
         # Extract and sort argument names excluding 'self'
         return sorted(p.name for p in parameters)
 
-    def get_params(self, deep: bool = True) -> Dict[str, Any]:
+    def get_params(self, deep: bool = True) -> dict[str, Any]:
         """
         Get the snowflake-ml parameters for this transformer.
 
@@ -255,7 +255,7 @@ class Base:
         Returns:
             Parameter names mapped to their values.
         """
-        out: Dict[str, Any] = dict()
+        out: dict[str, Any] = dict()
         for key in self._get_param_names():
             if hasattr(self, key):
                 value = getattr(self, key)
@@ -320,11 +320,11 @@ class Base:
         sklearn_initial_keywords: Optional[Union[str, Iterable[str]]] = None,
         sklearn_unused_keywords: Optional[Union[str, Iterable[str]]] = None,
         snowml_only_keywords: Optional[Union[str, Iterable[str]]] = None,
-        sklearn_added_keyword_to_version_dict: Optional[Dict[str, str]] = None,
-        sklearn_added_kwarg_value_to_version_dict: Optional[Dict[str, Dict[str, str]]] = None,
-        sklearn_deprecated_keyword_to_version_dict: Optional[Dict[str, str]] = None,
-        sklearn_removed_keyword_to_version_dict: Optional[Dict[str, str]] = None,
-    ) -> Dict[str, Any]:
+        sklearn_added_keyword_to_version_dict: Optional[dict[str, str]] = None,
+        sklearn_added_kwarg_value_to_version_dict: Optional[dict[str, dict[str, str]]] = None,
+        sklearn_deprecated_keyword_to_version_dict: Optional[dict[str, str]] = None,
+        sklearn_removed_keyword_to_version_dict: Optional[dict[str, str]] = None,
+    ) -> dict[str, Any]:
         """
         Get sklearn keyword arguments.
 
@@ -350,7 +350,7 @@ class Base:
         """
         default_sklearn_args = _utils.get_default_args(default_sklearn_obj.__class__.__init__)
         given_args = self.get_params()
-        sklearn_args: Dict[str, Any] = _utils.get_filtered_valid_sklearn_args(
+        sklearn_args: dict[str, Any] = _utils.get_filtered_valid_sklearn_args(
             args=given_args,
             default_sklearn_args=default_sklearn_args,
             sklearn_initial_keywords=sklearn_initial_keywords,
@@ -368,8 +368,8 @@ class BaseEstimator(Base):
     def __init__(
         self,
         *,
-        file_names: Optional[List[str]] = None,
-        custom_states: Optional[List[str]] = None,
+        file_names: Optional[list[str]] = None,
+        custom_states: Optional[list[str]] = None,
         sample_weight_col: Optional[str] = None,
     ) -> None:
         """
@@ -418,7 +418,7 @@ class BaseEstimator(Base):
         self.sample_weight_col = sample_weight_col
         return self
 
-    def _get_dependencies(self) -> List[str]:
+    def _get_dependencies(self) -> list[str]:
         """
         Return the list of conda dependencies required to work with the object.
 
@@ -458,8 +458,8 @@ class BaseEstimator(Base):
         return dataset[self.input_cols]
 
     def _compute(
-        self, dataset: snowpark.DataFrame, cols: List[str], states: List[str]
-    ) -> Dict[str, Dict[str, Union[int, float, str]]]:
+        self, dataset: snowpark.DataFrame, cols: list[str], states: list[str]
+    ) -> dict[str, dict[str, Union[int, float, str]]]:
         """
         Compute required states of the columns.
 
@@ -474,7 +474,7 @@ class BaseEstimator(Base):
             A dict of {column_name: {state: value}} of each column.
         """
 
-        def _compute_on_partition(df: snowpark.DataFrame, cols_subset: List[str]) -> snowpark.DataFrame:
+        def _compute_on_partition(df: snowpark.DataFrame, cols_subset: list[str]) -> snowpark.DataFrame:
             """Returns a DataFrame with the desired computation on the specified column subset."""
             exprs = []
             sql_prefix = "SQL>>>"
@@ -499,7 +499,7 @@ class BaseEstimator(Base):
             statement_params=telemetry.get_statement_params(PROJECT, SUBPROJECT, self.__class__.__name__),
         )
 
-        computed_dict: Dict[str, Dict[str, Union[int, float, str]]] = {}
+        computed_dict: dict[str, dict[str, Union[int, float, str]]] = {}
         for idx, val in enumerate(_results[0]):
             col_name = cols[idx // len(states)]
             if col_name not in computed_dict:
@@ -516,8 +516,8 @@ class BaseTransformer(BaseEstimator):
         self,
         *,
         drop_input_cols: Optional[bool] = False,
-        file_names: Optional[List[str]] = None,
-        custom_states: Optional[List[str]] = None,
+        file_names: Optional[list[str]] = None,
+        custom_states: Optional[list[str]] = None,
         sample_weight_col: Optional[str] = None,
     ) -> None:
         """Base class for all transformers."""
@@ -551,7 +551,7 @@ class BaseTransformer(BaseEstimator):
                 ),
             )
 
-    def _infer_input_cols(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> List[str]:
+    def _infer_input_cols(self, dataset: Union[snowpark.DataFrame, pd.DataFrame]) -> list[str]:
         """
         Infer input_cols from the dataset. Input column are all columns in the input dataset that are not
         designated as label, passthrough, or sample weight columns.
@@ -569,7 +569,7 @@ class BaseTransformer(BaseEstimator):
         ]
         return cols
 
-    def _infer_output_cols(self) -> List[str]:
+    def _infer_output_cols(self) -> list[str]:
         """Infer output column names from based on the estimator.
 
         Returns:
@@ -624,7 +624,7 @@ class BaseTransformer(BaseEstimator):
             cols = self._infer_output_cols()
             self.set_output_cols(output_cols=cols)
 
-    def _get_output_column_names(self, output_cols_prefix: str, output_cols: Optional[List[str]] = None) -> List[str]:
+    def _get_output_column_names(self, output_cols_prefix: str, output_cols: Optional[list[str]] = None) -> list[str]:
         """Returns the list of output columns for predict_proba(), decision_function(), etc.. functions.
         Returns a list with output_cols_prefix as the only element if the estimator is not a classifier.
 

snowflake/ml/modeling/impute/__init__.py

@@ -2,7 +2,7 @@ import os
 
 from snowflake.ml._internal import init_utils
 
-pkg_dir = os.path.dirname(os.path.abspath(__file__))
+pkg_dir = os.path.dirname(__file__)
 pkg_name = __name__
 exportable_classes = init_utils.fetch_classes_from_modules_in_pkg_dir(pkg_dir=pkg_dir, pkg_name=pkg_name)
 for k, v in exportable_classes.items():

snowflake/ml/modeling/impute/simple_imputer.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import copy
 import warnings
-from typing import Any, Dict, Iterable, Optional, Type, Union
+from typing import Any, Iterable, Optional, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -25,7 +25,7 @@ STRATEGY_TO_STATE_DICT = {
     "most_frequent": _utils.BasicStatistics.MODE,
 }
 
-SNOWFLAKE_DATATYPE_TO_NUMPY_DTYPE_MAP: Dict[Type[T.DataType], npt.DTypeLike] = {
+SNOWFLAKE_DATATYPE_TO_NUMPY_DTYPE_MAP: dict[type[T.DataType], npt.DTypeLike] = {
     T.ByteType: np.dtype("int8"),
     T.ShortType: np.dtype("int16"),
     T.IntegerType: np.dtype("int32"),
@@ -164,7 +164,7 @@ class SimpleImputer(base.BaseTransformer):
 
         self.fill_value = fill_value
         self.missing_values = missing_values
-        self.statistics_: Dict[str, Any] = {}
+        self.statistics_: dict[str, Any] = {}
         # TODO(hayu): [SNOW-752265] Support SimpleImputer keep_empty_features.
         #  Add back when `keep_empty_features` is supported.
         # self.keep_empty_features = keep_empty_features
@@ -195,7 +195,7 @@ class SimpleImputer(base.BaseTransformer):
             del self.feature_names_in_
             del self._sklearn_fit_dtype
 
-    def _get_dataset_input_col_datatypes(self, dataset: snowpark.DataFrame) -> Dict[str, T.DataType]:
+    def _get_dataset_input_col_datatypes(self, dataset: snowpark.DataFrame) -> dict[str, T.DataType]:
         """
         Checks that the input columns are all the same datatype category(except for most_frequent strategy) and
         returns the datatype.
@@ -211,7 +211,7 @@ class SimpleImputer(base.BaseTransformer):
                 supported.
         """
 
-        def check_type_consistency(col_types: Dict[str, T.DataType]) -> None:
+        def check_type_consistency(col_types: dict[str, T.DataType]) -> None:
             is_numeric_type = None
             for col_name, col_type in col_types.items():
                 if is_numeric_type is None:

snowflake/ml/modeling/metrics/__init__.py

@@ -5,7 +5,7 @@ import cloudpickle
 from snowflake.ml._internal import init_utils
 from snowflake.ml._internal.utils import result
 
-pkg_dir = os.path.dirname(os.path.abspath(__file__))
+pkg_dir = os.path.dirname(__file__)
 pkg_name = __name__
 exportable_functions = init_utils.fetch_functions_from_modules_in_pkg_dir(pkg_dir=pkg_dir, pkg_name=pkg_name)
 for k, v in exportable_functions.items():