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

snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Any, List, Optional
+from typing import Any, Optional
 
 import pandas as pd
 
@@ -38,9 +38,9 @@ class PandasTransformHandlers:
     def batch_inference(
         self,
         inference_method: str,
-        input_cols: List[str],
-        expected_output_cols: List[str],
-        snowpark_input_cols: Optional[List[str]] = None,
+        input_cols: list[str],
+        expected_output_cols: list[str],
+        snowpark_input_cols: Optional[list[str]] = None,
         drop_input_cols: Optional[bool] = False,
         *args: Any,
         **kwargs: Any,
@@ -147,8 +147,8 @@ class PandasTransformHandlers:
 
     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,

snowflake/ml/modeling/_internal/local_implementations/pandas_trainer.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import List, Optional, Tuple
+from typing import Optional
 
 import pandas as pd
 
@@ -15,8 +15,8 @@ class PandasModelTrainer:
         self,
         estimator: object,
         dataset: pd.DataFrame,
-        input_cols: List[str],
-        label_cols: Optional[List[str]],
+        input_cols: list[str],
+        label_cols: Optional[list[str]],
         sample_weight_col: Optional[str],
     ) -> None:
         """
@@ -57,10 +57,10 @@ class PandasModelTrainer:
 
     def train_fit_predict(
         self,
-        expected_output_cols_list: List[str],
+        expected_output_cols_list: list[str],
         drop_input_cols: Optional[bool] = False,
         example_output_pd_df: Optional[pd.DataFrame] = None,
-    ) -> Tuple[pd.DataFrame, object]:
+    ) -> tuple[pd.DataFrame, object]:
         """Trains the model using specified features and target columns from the dataset.
         This API is different from fit itself because it would also provide the predict
         output.
@@ -92,9 +92,9 @@ class PandasModelTrainer:
 
     def train_fit_transform(
         self,
-        expected_output_cols_list: List[str],
+        expected_output_cols_list: list[str],
         drop_input_cols: Optional[bool] = False,
-    ) -> Tuple[pd.DataFrame, object]:
+    ) -> tuple[pd.DataFrame, object]:
         """Trains the model using specified features and target columns from the dataset.
         This API is different from fit itself because it would also provide the transform
         output.

snowflake/ml/modeling/_internal/model_specifications.py

@@ -1,5 +1,3 @@
-from typing import List
-
 import cloudpickle as cp
 import numpy as np
 
@@ -11,7 +9,7 @@ class ModelSpecifications:
     A dataclass to define model based specifications like required imports, and package dependencies for Sproc/Udfs.
     """
 
-    def __init__(self, imports: List[str], pkgDependencies: List[str]) -> None:
+    def __init__(self, imports: list[str], pkgDependencies: list[str]) -> None:
         self.imports = imports
         self.pkgDependencies = pkgDependencies
 
@@ -20,7 +18,7 @@ class SKLearnModelSpecifications(ModelSpecifications):
     def __init__(self) -> None:
         import sklearn
 
-        imports: List[str] = ["sklearn"]
+        imports: list[str] = ["sklearn"]
         # TODO(snandamuri): Replace cloudpickle with joblib after latest version of joblib is added to snowflake conda.
         pkgDependencies = [
             f"numpy=={np.__version__}",
@@ -56,8 +54,8 @@ class XGBoostModelSpecifications(ModelSpecifications):
         import sklearn
         import xgboost
 
-        imports: List[str] = ["xgboost"]
-        pkgDependencies: List[str] = [
+        imports: list[str] = ["xgboost"]
+        pkgDependencies: list[str] = [
             f"numpy=={np.__version__}",
             f"scikit-learn=={sklearn.__version__}",
             f"xgboost=={xgboost.__version__}",
@@ -71,8 +69,8 @@ class LightGBMModelSpecifications(ModelSpecifications):
         import lightgbm
         import sklearn
 
-        imports: List[str] = ["lightgbm"]
-        pkgDependencies: List[str] = [
+        imports: list[str] = ["lightgbm"]
+        pkgDependencies: list[str] = [
             f"numpy=={np.__version__}",
             f"scikit-learn=={sklearn.__version__}",
             f"lightgbm=={lightgbm.__version__}",
@@ -86,8 +84,8 @@ class SklearnModelSelectionModelSpecifications(ModelSpecifications):
         import sklearn
         import xgboost
 
-        imports: List[str] = ["sklearn", "xgboost"]
-        pkgDependencies: List[str] = [
+        imports: list[str] = ["sklearn", "xgboost"]
+        pkgDependencies: list[str] = [
             f"numpy=={np.__version__}",
             f"scikit-learn=={sklearn.__version__}",
             f"cloudpickle=={cp.__version__}",

snowflake/ml/modeling/_internal/model_trainer.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Protocol, Tuple, Union
+from typing import Optional, Protocol, Union
 
 import pandas as pd
 
@@ -18,15 +18,15 @@ class ModelTrainer(Protocol):
 
     def train_fit_predict(
         self,
-        expected_output_cols_list: List[str],
+        expected_output_cols_list: list[str],
         drop_input_cols: Optional[bool] = False,
         example_output_pd_df: Optional[pd.DataFrame] = None,
-    ) -> Tuple[Union[DataFrame, pd.DataFrame], object]:
+    ) -> tuple[Union[DataFrame, pd.DataFrame], object]:
         raise NotImplementedError
 
     def train_fit_transform(
         self,
-        expected_output_cols_list: List[str],
+        expected_output_cols_list: list[str],
         drop_input_cols: Optional[bool] = False,
-    ) -> Tuple[Union[DataFrame, pd.DataFrame], object]:
+    ) -> tuple[Union[DataFrame, pd.DataFrame], object]:
         raise NotImplementedError

snowflake/ml/modeling/_internal/model_trainer_builder.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import Optional, Union
 
 import pandas as pd
 from sklearn import model_selection
@@ -71,8 +71,8 @@ class ModelTrainerBuilder:
         cls,
         estimator: object,
         dataset: Union[DataFrame, pd.DataFrame],
-        input_cols: Optional[List[str]] = None,
-        label_cols: Optional[List[str]] = None,
+        input_cols: Optional[list[str]] = None,
+        label_cols: Optional[list[str]] = None,
         sample_weight_col: Optional[str] = None,
         autogenerated: bool = False,
         subproject: str = "",
@@ -130,7 +130,7 @@ class ModelTrainerBuilder:
         cls,
         estimator: object,
         dataset: Union[DataFrame, pd.DataFrame],
-        input_cols: List[str],
+        input_cols: list[str],
         autogenerated: bool = False,
         subproject: str = "",
     ) -> ModelTrainer:
@@ -169,8 +169,8 @@ class ModelTrainerBuilder:
         cls,
         estimator: object,
         dataset: Union[DataFrame, pd.DataFrame],
-        input_cols: List[str],
-        label_cols: Optional[List[str]] = None,
+        input_cols: list[str],
+        label_cols: Optional[list[str]] = None,
         sample_weight_col: Optional[str] = None,
         autogenerated: bool = False,
         subproject: str = "",

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

@@ -5,7 +5,7 @@ import os
 import posixpath
 import sys
 import uuid
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional, Union
 
 import cloudpickle as cp
 import numpy as np
@@ -50,11 +50,11 @@ _UDTF_STAGE_NAME = f"MEMORY_EFFICIENT_UDTF_{str(uuid.uuid4()).replace('-', '_')}
 def construct_cv_results(
     estimator: Union[GridSearchCV, RandomizedSearchCV],
     n_split: int,
-    param_grid: List[Dict[str, Any]],
-    cv_results_raw_hex: List[Row],
+    param_grid: list[dict[str, Any]],
+    cv_results_raw_hex: list[Row],
     cross_validator_indices_length: int,
     parameter_grid_length: int,
-) -> Tuple[bool, Dict[str, Any]]:
+) -> tuple[bool, dict[str, Any]]:
     """Construct the cross validation result from the UDF. Because we accelerate the process
     by the number of cross validation number, and the combination of parameter grids.
     Therefore, we need to stick them back together instead of returning the raw result
@@ -158,11 +158,11 @@ def construct_cv_results(
 def construct_cv_results_memory_efficient_version(
     estimator: Union[GridSearchCV, RandomizedSearchCV],
     n_split: int,
-    param_grid: List[Dict[str, Any]],
-    cv_results_raw_hex: List[Row],
+    param_grid: list[dict[str, Any]],
+    cv_results_raw_hex: list[Row],
     cross_validator_indices_length: int,
     parameter_grid_length: int,
-) -> Tuple[Any, Dict[str, Any]]:
+) -> tuple[Any, dict[str, Any]]:
     """Construct the cross validation result from the UDF.
     The output is a raw dictionary generated by _fit_and_score, encoded into hex binary.
     This function need to decode the string and then call _format_result to stick them back together
@@ -210,7 +210,7 @@ def construct_cv_results_memory_efficient_version(
     # because original SearchCV is ranked by parameter first and cv second,
     # to make the memory efficient, we implemented by fitting on cv first and parameter second
     # when retrieving the results back, the ordering should revert back to remain the same result as original SearchCV
-    def generate_the_order_by_parameter_index(all_combination_length: int) -> List[int]:
+    def generate_the_order_by_parameter_index(all_combination_length: int) -> list[int]:
         pattern = []
         for i in range(all_combination_length):
             if i % parameter_grid_length == 0:
@@ -221,7 +221,7 @@ def construct_cv_results_memory_efficient_version(
                     pattern.append(j)
         return pattern
 
-    def rerank_array(original_array: List[Any], pattern: List[int]) -> List[Any]:
+    def rerank_array(original_array: list[Any], pattern: list[int]) -> list[Any]:
         reranked_array = []
         for index in pattern:
             reranked_array.append(original_array[index])
@@ -251,8 +251,8 @@ class DistributedHPOTrainer(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 = "",
@@ -289,10 +289,10 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
         dataset: DataFrame,
         session: Session,
         estimator: Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV],
-        dependencies: List[str],
-        udf_imports: List[str],
-        input_cols: List[str],
-        label_cols: Optional[List[str]],
+        dependencies: list[str],
+        udf_imports: list[str],
+        input_cols: list[str],
+        label_cols: Optional[list[str]],
         sample_weight_col: Optional[str],
     ) -> Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV]:
         from itertools import product
@@ -382,10 +382,10 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
         )
         def _distributed_search(
             session: Session,
-            imports: List[str],
+            imports: list[str],
             stage_estimator_file_name: str,
-            input_cols: List[str],
-            label_cols: Optional[List[str]],
+            input_cols: list[str],
+            label_cols: Optional[list[str]],
         ) -> str:
             import os
             import time
@@ -455,12 +455,12 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
             assert estimator is not None
 
             @cachetools.cached(cache={})
-            def _load_data_into_udf() -> Tuple[
-                Dict[str, pd.DataFrame],
+            def _load_data_into_udf() -> tuple[
+                dict[str, pd.DataFrame],
                 Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV],
                 pd.DataFrame,
                 int,
-                List[Dict[str, Any]],
+                list[dict[str, Any]],
             ]:
                 import pyarrow.parquet as pq
 
@@ -512,7 +512,7 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
                     self.data_length = data_length
                     self.params_to_evaluate = params_to_evaluate
 
-                def process(self, params_idx: int, cv_idx: int) -> Iterator[Tuple[str]]:
+                def process(self, params_idx: int, cv_idx: int) -> Iterator[tuple[str]]:
                     # Assign parameter to GridSearchCV
                     if hasattr(estimator, "param_grid"):
                         self.estimator.param_grid = self.params_to_evaluate[params_idx]
@@ -699,10 +699,10 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
         dataset: DataFrame,
         session: Session,
         estimator: Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV],
-        dependencies: List[str],
-        udf_imports: List[str],
-        input_cols: List[str],
-        label_cols: Optional[List[str]],
+        dependencies: list[str],
+        udf_imports: list[str],
+        input_cols: list[str],
+        label_cols: Optional[list[str]],
         sample_weight_col: Optional[str],
     ) -> Union[model_selection.GridSearchCV, model_selection.RandomizedSearchCV]:
         from itertools import product
@@ -727,7 +727,7 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
         # Create a temp file and dump the estimator to that file.
         estimator_file_name = temp_file_utils.get_temp_file_path()
         params_to_evaluate = list(param_grid)
-        CONSTANTS: Dict[str, Any] = dict()
+        CONSTANTS: dict[str, Any] = dict()
         CONSTANTS["dataset_snowpark_cols"] = dataset.columns
         CONSTANTS["n_candidates"] = len(params_to_evaluate)
         CONSTANTS["_N_JOBS"] = estimator.n_jobs
@@ -791,10 +791,10 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
         )
         def _distributed_search(
             session: Session,
-            imports: List[str],
+            imports: list[str],
             stage_estimator_file_name: str,
-            input_cols: List[str],
-            label_cols: Optional[List[str]],
+            input_cols: list[str],
+            label_cols: Optional[list[str]],
         ) -> str:
             import os
             import time

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

@@ -3,7 +3,7 @@ import inspect
 import os
 import posixpath
 import sys
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 from uuid import uuid4
 
 import cloudpickle as cp
@@ -73,10 +73,10 @@ class SnowparkTransformHandlers:
     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: Session,
-        dependencies: List[str],
+        dependencies: list[str],
         drop_input_cols: Optional[bool] = False,
         expected_output_cols_type: Optional[str] = "",
         *args: Any,
@@ -229,11 +229,11 @@ class SnowparkTransformHandlers:
 
     def score(
         self,
-        input_cols: List[str],
-        label_cols: List[str],
+        input_cols: list[str],
+        label_cols: list[str],
         session: 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,
@@ -308,12 +308,12 @@ class SnowparkTransformHandlers:
         )
         def score_wrapper_sproc(
             session: Session,
-            sql_queries: List[str],
+            sql_queries: list[str],
             stage_score_file_name: str,
-            input_cols: List[str],
-            label_cols: List[str],
+            input_cols: list[str],
+            label_cols: list[str],
             sample_weight_col: Optional[str],
-            score_statement_params: Dict[str, str],
+            score_statement_params: dict[str, str],
         ) -> float:
             import inspect
             import os
@@ -382,7 +382,7 @@ class SnowparkTransformHandlers:
 
         return score
 
-    def _get_validated_snowpark_dependencies(self, session: Session, dependencies: List[str]) -> List[str]:
+    def _get_validated_snowpark_dependencies(self, session: Session, dependencies: list[str]) -> list[str]:
         """A helper function to validate dependencies and return the available packages that exists
         in the snowflake anaconda channel
 

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

@@ -2,7 +2,7 @@ import importlib
 import inspect
 import os
 import posixpath
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Union
 
 import cloudpickle as cp
 import pandas as pd
@@ -55,8 +55,8 @@ class 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 = "",
@@ -84,7 +84,7 @@ class SnowparkModelTrainer:
         self._subproject = subproject
         self._class_name = estimator.__class__.__name__
 
-    def _fetch_model_from_stage(self, dir_path: str, file_name: str, statement_params: Dict[str, str]) -> object:
+    def _fetch_model_from_stage(self, dir_path: str, file_name: str, statement_params: dict[str, str]) -> object:
         """
         Downloads the serialized model from a stage location and unpickles it.
 
@@ -112,7 +112,7 @@ class SnowparkModelTrainer:
     def _build_fit_wrapper_sproc(
         self,
         model_spec: ModelSpecifications,
-    ) -> Callable[[Any, List[str], str, List[str], List[str], Optional[str], Dict[str, str]], str]:
+    ) -> Callable[[Any, list[str], str, list[str], list[str], Optional[str], dict[str, str]], str]:
         """
         Constructs and returns a python stored procedure function to be used for training model.
 
@@ -129,12 +129,12 @@ class SnowparkModelTrainer:
 
         def fit_wrapper_function(
             session: Session,
-            sql_queries: List[str],
+            sql_queries: list[str],
             temp_stage_name: str,
-            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 inspect
             import os
@@ -218,7 +218,7 @@ class SnowparkModelTrainer:
 
         return fit_wrapper_function
 
-    def _get_fit_wrapper_sproc(self, statement_params: Dict[str, str], anonymous: bool) -> StoredProcedure:
+    def _get_fit_wrapper_sproc(self, statement_params: dict[str, str], anonymous: bool) -> StoredProcedure:
         model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
 
         fit_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
@@ -243,7 +243,7 @@ class SnowparkModelTrainer:
     def _build_fit_predict_wrapper_sproc(
         self,
         model_spec: ModelSpecifications,
-    ) -> Callable[[Session, List[str], str, List[str], Dict[str, str], bool, List[str], str], str]:
+    ) -> Callable[[Session, list[str], str, list[str], dict[str, str], bool, list[str], str], str]:
         """
         Constructs and returns a python stored procedure function to be used for training model.
 
@@ -258,12 +258,12 @@ class SnowparkModelTrainer:
 
         def fit_predict_wrapper_function(
             session: Session,
-            sql_queries: List[str],
+            sql_queries: list[str],
             temp_stage_name: str,
-            input_cols: List[str],
-            statement_params: Dict[str, str],
+            input_cols: list[str],
+            statement_params: dict[str, str],
             drop_input_cols: bool,
-            expected_output_cols_list: List[str],
+            expected_output_cols_list: list[str],
             fit_predict_result_name: str,
         ) -> str:
             import os
@@ -346,14 +346,14 @@ class SnowparkModelTrainer:
     ) -> Callable[
         [
             Session,
-            List[str],
+            list[str],
             str,
-            List[str],
-            Optional[List[str]],
+            list[str],
+            Optional[list[str]],
             Optional[str],
-            Dict[str, str],
+            dict[str, str],
             bool,
-            List[str],
+            list[str],
             str,
         ],
         str,
@@ -372,14 +372,14 @@ class SnowparkModelTrainer:
 
         def fit_transform_wrapper_function(
             session: Session,
-            sql_queries: List[str],
+            sql_queries: list[str],
             temp_stage_name: str,
-            input_cols: List[str],
-            label_cols: Optional[List[str]],
+            input_cols: list[str],
+            label_cols: Optional[list[str]],
             sample_weight_col: Optional[str],
-            statement_params: Dict[str, str],
+            statement_params: dict[str, str],
             drop_input_cols: bool,
-            expected_output_cols_list: List[str],
+            expected_output_cols_list: list[str],
             fit_transform_result_name: str,
         ) -> str:
             import os
@@ -473,7 +473,7 @@ class SnowparkModelTrainer:
 
         return fit_transform_wrapper_function
 
-    def _get_fit_predict_wrapper_sproc(self, statement_params: Dict[str, str], anonymous: bool) -> StoredProcedure:
+    def _get_fit_predict_wrapper_sproc(self, statement_params: dict[str, str], anonymous: bool) -> StoredProcedure:
         model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
 
         fit_predict_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
@@ -495,7 +495,7 @@ class SnowparkModelTrainer:
 
         return fit_predict_wrapper_sproc
 
-    def _get_fit_transform_wrapper_sproc(self, statement_params: Dict[str, str], anonymous: bool) -> StoredProcedure:
+    def _get_fit_transform_wrapper_sproc(self, statement_params: dict[str, str], anonymous: bool) -> StoredProcedure:
         model_spec = ModelSpecificationsBuilder.build(model=self.estimator)
 
         fit_transform_sproc_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.PROCEDURE)
@@ -586,10 +586,10 @@ class SnowparkModelTrainer:
 
     def train_fit_predict(
         self,
-        expected_output_cols_list: List[str],
+        expected_output_cols_list: list[str],
         drop_input_cols: Optional[bool] = False,
         example_output_pd_df: Optional[pd.DataFrame] = None,
-    ) -> Tuple[Union[DataFrame, pd.DataFrame], object]:
+    ) -> tuple[Union[DataFrame, pd.DataFrame], object]:
         """Trains the model by pushing down the compute into Snowflake using stored procedures.
         This API is different from fit itself because it would also provide the predict
         output.
@@ -682,9 +682,9 @@ class SnowparkModelTrainer:
 
     def train_fit_transform(
         self,
-        expected_output_cols_list: List[str],
+        expected_output_cols_list: list[str],
         drop_input_cols: Optional[bool] = False,
-    ) -> Tuple[Union[DataFrame, pd.DataFrame], object]:
+    ) -> tuple[Union[DataFrame, pd.DataFrame], object]:
         """Trains the model by pushing down the compute into Snowflake using stored procedures.
         This API is different from fit itself because it would also provide the transform
         output.