sklearn converted to numpy

georgia-tech-db · Nov 10, 2023 · 823e3fe · 823e3fe
1 parent 1d7c30e
commit 823e3fe
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/evadb/executor/create_function_executor.py b/evadb/executor/create_function_executor.py
@@ -24,7 +24,6 @@
 
 import numpy as np
 import pandas as pd
-from sklearn.metrics import mean_squared_error
 
 from evadb.catalog.catalog_utils import get_metadata_properties
 from evadb.catalog.models.function_catalog import FunctionCatalogEntry
@@ -57,6 +56,10 @@
 from evadb.utils.logging_manager import logger
 
 
+def root_mean_squared_error(y_true, y_pred):
+    return np.sqrt(np.mean(np.square(y_pred - y_true)))
+
+
 # From https://stackoverflow.com/a/34333710
 @contextlib.contextmanager
 def set_env(**environ):
@@ -602,12 +605,11 @@ def get_optuna_config(trial):
                                 crossvalidation_df.unique_id == uid
                             ]
                             rmses.append(
-                                mean_squared_error(
+                                root_mean_squared_error(
                                     crossvalidation_df_here.y,
                                     crossvalidation_df_here[
                                         arg_map["model"] + "-median"
                                     ],
-                                    squared=False,
                                 )
                                 / np.mean(crossvalidation_df_here.y)
                             )
@@ -643,10 +645,9 @@ def get_optuna_config(trial):
                             crossvalidation_df.unique_id == uid
                         ]
                         rmses.append(
-                            mean_squared_error(
+                            root_mean_squared_error(
                                 crossvalidation_df_here.y,
                                 crossvalidation_df_here[arg_map["model"]],
-                                squared=False,
                             )
                             / np.mean(crossvalidation_df_here.y)
                         )