Overload model_fit_params instead

kingychiu · Nov 13, 2023 · 0d34aa8 · 0d34aa8
1 parent 3c341e2
commit 0d34aa8
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 18 deletions.
diff --git a/benchmarks/run_tabular_benchmark.py b/benchmarks/run_tabular_benchmark.py
@@ -239,7 +239,7 @@ def write_report(reports: List):
         importance_dfs = compute(
             model_cls=model_cls,
             model_cls_params=model_cls_params,
-            model_fit_params_builder=lambda _: {},
+            model_fit_params=lambda _: {},
             X=X_train,
             y=y_train,
             num_actual_runs=num_actual_runs,

diff --git a/target_permutation_importances/functional.py b/target_permutation_importances/functional.py
@@ -329,7 +329,7 @@ def _get_model_importances_attr(model: Any):
 def compute(
     model_cls: Any,
     model_cls_params: Dict,
-    model_fit_params_builder: ModelFitParamsBuilderType,
+    model_fit_params: Union[ModelFitParamsBuilderType, Dict],
     X: XType,
     y: YType,
     num_actual_runs: PositiveInt = 2,
@@ -345,7 +345,7 @@ def compute(
     Args:
         model_cls: The constructor/class of the model.
         model_cls_params: The parameters to pass to the model constructor.
-        model_fit_params_builder: A function that return parameters to pass to the model fit method.
+        model_fit_params: A function that return parameters to pass to the model fit method.
         X: The input data.
         y: The target vector.
         num_actual_runs: Number of actual runs. Defaults to 2.
@@ -379,7 +379,7 @@ def compute(
             model_cls_params={ # The parameters to pass to the model constructor. Update this based on your needs.
                 "n_jobs": -1,
             },
-            model_fit_params_builder=lambda _: {}, # The parameters to pass to the model fit method. Update this based on your needs.
+            model_fit_params=lambda _: {}, # The parameters to pass to the model fit method. Update this based on your needs.
             X=Xpd, # pd.DataFrame, np.ndarray
             y=data.target, # pd.Series, np.ndarray
             num_actual_runs=2,
@@ -424,9 +424,12 @@ def _model_builder(is_random_run: bool, run_idx: int) -> Any:
         return model_cls(**_model_cls_params)
 
     def _model_fitter(model: Any, X: XType, y: YType) -> Any:
-        _model_fit_params = model_fit_params_builder(
-            list(X.columns) if isinstance(X, pd.DataFrame) else None,
-        )
+        if isinstance(model_fit_params, dict):
+            _model_fit_params = model_fit_params.copy()
+        else:
+            _model_fit_params = model_fit_params(
+                list(X.columns) if isinstance(X, pd.DataFrame) else None,
+            )
         if "Cat" in str(model.__class__):
             _model_fit_params["verbose"] = False
         return model.fit(X, y, **_model_fit_params)

diff --git a/target_permutation_importances/sklearn_wrapper.py b/target_permutation_importances/sklearn_wrapper.py
@@ -139,7 +139,7 @@ def fit(
         result = compute(
             model_cls=self.model_cls,
             model_cls_params=self.model_cls_params,
-            model_fit_params_builder=lambda _: fit_params,
+            model_fit_params=lambda _: fit_params,
             X=X,
             y=y,
             num_actual_runs=self.num_actual_runs,

diff --git a/tests/test_functional.py b/tests/test_functional.py
@@ -70,7 +70,7 @@ def test_compute_binary_classification(model_cls, imp_func, xtype):
     result_df = compute(
         model_cls=model_cls[0],
         model_cls_params=model_cls[1],
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=imp_func,
         X=X,
         y=data.target,
@@ -103,7 +103,7 @@ def test_compute_multi_class_classification(model_cls, imp_func, xtype):
     result_df = compute(
         model_cls=model_cls[0],
         model_cls_params=model_cls[1],
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=imp_func,
         X=X,
         y=data.target,
@@ -146,7 +146,7 @@ def test_compute_multi_label_classification(model_cls, imp_func, xtype):
     result_df = compute(
         model_cls=model_cls[0],
         model_cls_params=model_cls_params,
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=imp_func,
         X=X,
         y=y,
@@ -185,7 +185,7 @@ def test_compute_multi_label_classification_with_MultiOutputClassifier(
         model_cls_params={
             "estimator": model_cls[0](**model_cls[1]),
         },
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=imp_func,
         X=X,
         y=y,
@@ -216,7 +216,7 @@ def test_compute_regression(model_cls, imp_func, xtype):
     result_df = compute(
         model_cls=model_cls[0],
         model_cls_params=model_cls[1],
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=imp_func,
         X=X,
         y=data.target,
@@ -255,7 +255,7 @@ def test_compute_multi_target_regression_with_MultiOutputRegressor(
         model_cls_params={
             "estimator": model_cls[0](**model_cls[1]),
         },
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=imp_func,
         X=X,
         y=y,
@@ -282,7 +282,7 @@ def test_compute_with_multiple_importance_functions():
     result_dfs = compute(
         model_cls=RandomForestClassifier,
         model_cls_params={"n_estimators": 2, "n_jobs": 1},
-        model_fit_params_builder=lambda _: {},
+        model_fit_params=lambda _: {},
         permutation_importance_calculator=[
             compute_permutation_importance_by_subtraction,
             compute_permutation_importance_by_division,
@@ -371,7 +371,7 @@ def test_invalid_compute():
         compute(
             model_cls=RandomForestClassifier,
             model_cls_params={},
-            model_fit_params_builder=lambda _: {},
+            model_fit_params=lambda _: {},
             permutation_importance_calculator=compute_permutation_importance_by_subtraction,
             X=1,
             y=data.target,
@@ -382,7 +382,7 @@ def test_invalid_compute():
         compute(
             model_cls=RandomForestClassifier,
             model_cls_params={},
-            model_fit_params_builder=lambda _: {},
+            model_fit_params=lambda _: {},
             permutation_importance_calculator=compute_permutation_importance_by_subtraction,
             X=Xpd,
             y=1,
@@ -393,7 +393,7 @@ def test_invalid_compute():
         compute(
             model_cls=RandomForestClassifier,
             model_cls_params={},
-            model_fit_params_builder=lambda _: {},
+            model_fit_params=lambda _: {},
             permutation_importance_calculator=compute_permutation_importance_by_subtraction,
             X=Xpd,
             y=data.target,