fix sklearn and xgboost test errors

georgia-tech-db · Nov 16, 2023 · 849d142 · 849d142
1 parent eff6f60
commit 849d142
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/evadb/binder/statement_binder.py b/evadb/binder/statement_binder.py
@@ -92,7 +92,9 @@ def _bind_create_function_statement(self, node: CreateFunctionStatement):
                         outputs.append(column)
                     else:
                         inputs.append(column)
-            elif string_comparison_case_insensitive(node.function_type, "sklearn"):
+            elif string_comparison_case_insensitive(
+                node.function_type, "sklearn"
+            ) or string_comparison_case_insensitive(node.function_type, "XGBoost"):
                 assert (
                     "predict" in arg_map
                 ), f"Creating {node.function_type} functions expects 'predict' metadata."

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