snowflake-ml-python 1.8.2__py3-none-any.whl → 1.8.4__py3-none-any.whl

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.

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/calibration/calibrated_classifier_cv.py

@@ -11,7 +11,7 @@ import cloudpickle as cp
 import numpy as np
 import pandas as pd
 from numpy import typing as npt
-
+from packaging import version
 
 import numpy
 import sklearn
@@ -60,6 +60,14 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.6')
+# Modeling library estimators require a smaller sklearn version range.
+if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
+    raise Exception(
+        f"To use the modeling library, install scikit-learn version >= {SKLEARN_LOWER} and < {SKLEARN_UPPER}"
+    )
+
+
 class CalibratedClassifierCV(BaseTransformer):
     r"""Probability calibration with isotonic regression or logistic regression
     For more details on this class, see [sklearn.calibration.CalibratedClassifierCV]

snowflake/ml/modeling/cluster/affinity_propagation.py

@@ -11,7 +11,7 @@ import cloudpickle as cp
 import numpy as np
 import pandas as pd
 from numpy import typing as npt
-
+from packaging import version
 
 import numpy
 import sklearn
@@ -60,6 +60,14 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.6')
+# Modeling library estimators require a smaller sklearn version range.
+if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
+    raise Exception(
+        f"To use the modeling library, install scikit-learn version >= {SKLEARN_LOWER} and < {SKLEARN_UPPER}"
+    )
+
+
 class AffinityPropagation(BaseTransformer):
     r"""Perform Affinity Propagation Clustering of data
     For more details on this class, see [sklearn.cluster.AffinityPropagation]

snowflake/ml/modeling/cluster/agglomerative_clustering.py

@@ -11,7 +11,7 @@ import cloudpickle as cp
 import numpy as np
 import pandas as pd
 from numpy import typing as npt
-
+from packaging import version
 
 import numpy
 import sklearn
@@ -60,6 +60,14 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.6')
+# Modeling library estimators require a smaller sklearn version range.
+if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
+    raise Exception(
+        f"To use the modeling library, install scikit-learn version >= {SKLEARN_LOWER} and < {SKLEARN_UPPER}"
+    )
+
+
 class AgglomerativeClustering(BaseTransformer):
     r"""Agglomerative Clustering
     For more details on this class, see [sklearn.cluster.AgglomerativeClustering]

snowflake/ml/modeling/cluster/birch.py

@@ -11,7 +11,7 @@ import cloudpickle as cp
 import numpy as np
 import pandas as pd
 from numpy import typing as npt
-
+from packaging import version
 
 import numpy
 import sklearn
@@ -60,6 +60,14 @@ DATAFRAME_TYPE = Union[DataFrame, pd.DataFrame]
 
 INFER_SIGNATURE_MAX_ROWS = 100
 
+SKLEARN_LOWER, SKLEARN_UPPER = ('1.4', '1.6')
+# Modeling library estimators require a smaller sklearn version range.
+if not version.Version(SKLEARN_LOWER) <= version.Version(sklearn.__version__) < version.Version(SKLEARN_UPPER):
+    raise Exception(
+        f"To use the modeling library, install scikit-learn version >= {SKLEARN_LOWER} and < {SKLEARN_UPPER}"
+    )
+
+
 class Birch(BaseTransformer):
     r"""Implements the BIRCH clustering algorithm
     For more details on this class, see [sklearn.cluster.Birch]