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

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/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():