fix(Tunable): updated interface

dream-faster · May 11, 2023 · 14fcea9 · 14fcea9
1 parent 29c60f5
commit 14fcea9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/fold_wrappers/lightgbm.py b/src/fold_wrappers/lightgbm.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Type, Union
+from typing import Any, Callable, Optional, Type, Union
 
 import pandas as pd
 from fold.base import Tunable
@@ -81,7 +81,9 @@ def predict(self, X: pd.DataFrame) -> Union[pd.Series, pd.DataFrame]:
     def get_params(self) -> dict:
         return self.model.get_params()
 
-    def clone_with_params(self, **parameters) -> Tunable:
+    def clone_with_params(
+        self, parameters: dict, clone_children: Optional[Callable] = None
+    ) -> Tunable:
         return WrapLGBM(
             self.model_class,
             init_args=parameters,

diff --git a/src/fold_wrappers/xgboost.py b/src/fold_wrappers/xgboost.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Type, Union
+from typing import Any, Callable, Optional, Type, Union
 
 import pandas as pd
 from fold.base import Tunable
@@ -85,7 +85,9 @@ def predict(self, X: pd.DataFrame) -> Union[pd.Series, pd.DataFrame]:
     def get_params(self) -> dict:
         return self.model.get_params()
 
-    def clone_with_params(self, **parameters) -> Tunable:
+    def clone_with_params(
+        self, parameters: dict, clone_children: Optional[Callable] = None
+    ) -> Tunable:
         return WrapXGB(
             model_class=self.model_class,
             init_args=parameters,