snowflake-ml-python 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

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

@@ -1,10 +1,10 @@
 import inspect
 from typing import Any, List, Optional
 
-import numpy as np
 import pandas as pd
 
 from snowflake.ml._internal.exceptions import error_codes, exceptions
+from snowflake.ml.modeling._internal.estimator_utils import handle_inference_result
 
 
 class PandasTransformHandlers:
@@ -107,48 +107,9 @@ class PandasTransformHandlers:
 
         inference_res = getattr(self.estimator, inference_method)(input_df, *args, **kwargs)
 
-        if isinstance(inference_res, list) and len(inference_res) > 0 and isinstance(inference_res[0], np.ndarray):
-            # In case of multioutput estimators, predict_proba, decision_function etc., functions return a list of
-            # ndarrays. We need to concatenate them.
-
-            # First compute output column names
-            if len(output_cols) == len(inference_res):
-                actual_output_cols = []
-                for idx, np_arr in enumerate(inference_res):
-                    for i in range(1 if len(np_arr.shape) <= 1 else np_arr.shape[1]):
-                        actual_output_cols.append(f"{output_cols[idx]}_{i}")
-                output_cols = actual_output_cols
-
-            # Concatenate np arrays
-            transformed_numpy_array = np.concatenate(inference_res, axis=1)
-        elif isinstance(inference_res, tuple) and len(inference_res) > 0 and isinstance(inference_res[0], np.ndarray):
-            # In case of kneighbors, functions return a tuple of ndarrays.
-            transformed_numpy_array = np.stack(inference_res, axis=1)
-        else:
-            transformed_numpy_array = inference_res
-
-        if (len(transformed_numpy_array.shape) == 3) and inference_method != "kneighbors":
-            # VotingClassifier will return results of shape (n_classifiers, n_samples, n_classes)
-            # when voting = "soft" and flatten_transform = False. We can't handle unflatten transforms,
-            # so we ignore flatten_transform flag and flatten the results.
-            transformed_numpy_array = np.hstack(transformed_numpy_array)  # type: ignore[call-overload]
-
-        if len(transformed_numpy_array.shape) == 1:
-            transformed_numpy_array = np.reshape(transformed_numpy_array, (-1, 1))
-
-        shape = transformed_numpy_array.shape
-        if shape[1] != len(output_cols):
-            if len(output_cols) != 1:
-                raise exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_ARGUMENT,
-                    original_exception=TypeError(
-                        "expected_output_cols must be same length as transformed array or " "should be of length 1"
-                    ),
-                )
-            actual_output_cols = []
-            for i in range(shape[1]):
-                actual_output_cols.append(f"{output_cols[0]}_{i}")
-            output_cols = actual_output_cols
+        transformed_numpy_array, output_cols = handle_inference_result(
+            inference_res=inference_res, output_cols=output_cols, inference_method=inference_method
+        )
 
         if inference_method == "kneighbors":
             if len(transformed_numpy_array.shape) == 3:  # return_distance=True

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

@@ -55,18 +55,18 @@ class PandasModelTrainer:
 
     def train_fit_predict(
         self,
-        pass_through_columns: List[str],
         expected_output_cols_list: List[str],
+        drop_input_cols: Optional[bool] = False,
     ) -> 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.
 
         Args:
-            pass_through_columns (List[str]): The column names that would
-                display in the returned dataset.
             expected_output_cols_list (List[str]): The output columns
                 name as a list. Defaults to None.
+            drop_input_cols (Optional[bool]): Boolean to determine whether to
+                drop the input columns from the output dataset.
 
         Returns:
             Tuple[pd.DataFrame, object]: [predicted dataset, estimator]
@@ -75,7 +75,7 @@ class PandasModelTrainer:
         args = {"X": self.dataset[self.input_cols]}
         result = self.estimator.fit_predict(**args)
         result_df = pd.DataFrame(data=result, columns=expected_output_cols_list)
-        if len(pass_through_columns) == 0:
+        if drop_input_cols:
             result_df = result_df
         else:
             result_df = pd.concat([self.dataset, result_df], axis=1)

snowflake/ml/modeling/_internal/ml_runtime_implementations/ml_runtime_handlers.py

@@ -1,5 +1,8 @@
 from typing import Any, List, Optional
 
+from snowflake.ml.modeling._internal.snowpark_implementations.snowpark_handlers import (
+    SnowparkTransformHandlers,
+)
 from snowflake.snowpark import DataFrame, Session
 
 
@@ -42,14 +45,16 @@ class MLRuntimeTransformHandlers:
         inference_method: str,
         input_cols: List[str],
         expected_output_cols: List[str],
-        pass_through_cols: List[str],
         session: Session,
         dependencies: List[str],
+        drop_input_cols: Optional[bool] = False,
         expected_output_cols_type: Optional[str] = "",
         *args: Any,
         **kwargs: Any,
     ) -> DataFrame:
         """Run batch inference on the given dataset.
+        Temporary workaround - pushdown implementation is not currently ready for batch_inference.
+        We use a SnowparkTransformHandlers until we have a way to use the runtime client.
 
         Args:
             inference_method: the name of the method used by `estimator` to run inference.
@@ -57,7 +62,7 @@ class MLRuntimeTransformHandlers:
             session: An active Snowpark Session.
             dependencies: List of dependencies for the transformer.
             expected_output_cols: column names (in order) of the output dataset.
-            pass_through_cols: columns in the dataset not used in inference.
+            drop_input_cols: Boolean to determine whether to drop the input columns from the output dataset.
             expected_output_cols_type: Expected type of the output columns.
             args: additional positional arguments.
             kwargs: additional keyword args.
@@ -65,27 +70,26 @@ class MLRuntimeTransformHandlers:
         Returns:
             A new dataset of the same type as the input dataset.
 
-        Raises:
-            TypeError: The ML Runtimes client returned a non-DataFrame result.
         """
-        output_df = self.client.batch_inference(
-            inference_method=inference_method,
+
+        handler = SnowparkTransformHandlers(
             dataset=self.dataset,
             estimator=self.estimator,
-            input_cols=input_cols,
-            expected_output_cols=expected_output_cols,
-            pass_through_cols=pass_through_cols,
-            session=session,
-            dependencies=dependencies,
-            expected_output_cols_type=expected_output_cols_type,
+            class_name=self._class_name,
+            subproject=self._subproject,
+            autogenerated=self._autogenerated,
+        )
+        return handler.batch_inference(
+            inference_method,
+            input_cols,
+            expected_output_cols,
+            session,
+            dependencies,
+            drop_input_cols,
+            expected_output_cols_type,
             *args,
             **kwargs,
         )
-        if not isinstance(output_df, DataFrame):
-            raise TypeError(
-                f"The ML Runtimes Client did not return a DataFrame a non-float value Returned type: {type(output_df)}"
-            )
-        return output_df
 
     def score(
         self,

snowflake/ml/modeling/_internal/model_specifications.py

@@ -29,7 +29,7 @@ class SKLearnModelSpecifications(ModelSpecifications):
         ]
 
         # A change from previous implementation.
-        # When reusing the Sprocs for all the fit() call in the session, the static dpendencies list should include
+        # When reusing the Sprocs for all the fit() call in the session, the static dependencies list should include
         # all the possible dependencies required during the lifetime.
 
         # Include XGBoost in the dependencies if it is installed.
@@ -67,10 +67,12 @@ class XGBoostModelSpecifications(ModelSpecifications):
 class LightGBMModelSpecifications(ModelSpecifications):
     def __init__(self) -> None:
         import lightgbm
+        import sklearn
 
         imports: List[str] = ["lightgbm"]
         pkgDependencies: List[str] = [
             f"numpy=={np.__version__}",
+            f"scikit-learn=={sklearn.__version__}",
             f"lightgbm=={lightgbm.__version__}",
             f"cloudpickle=={cp.__version__}",
         ]

snowflake/ml/modeling/_internal/model_trainer.py

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