prevent Dataset object from re-inferring categorical features (#330)

deepchecks · Dec 27, 2021 · c4af9e8 · c4af9e8
1 parent 4b917d3
commit c4af9e8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/deepchecks/checks/performance/segment_performance.py b/deepchecks/checks/performance/segment_performance.py
@@ -112,7 +112,7 @@ def run(self, dataset, model) -> CheckResult:
                 else:
                     score = scorer(model,
                                    Dataset(feature_2_df, features=dataset.features,
-                                           label_name=dataset.label_name, cat_features=[]))
+                                           label_name=dataset.label_name, cat_features=dataset.cat_features))
                 scores[i, j] = score
                 counts[i, j] = len(feature_2_df)
 

diff --git a/deepchecks/checks/performance/simple_model_comparison.py b/deepchecks/checks/performance/simple_model_comparison.py
@@ -142,7 +142,7 @@ def _find_score(self, train_ds: Dataset, test_ds: Dataset, task_type: ModelType,
         scorer = get_scorer_single(model, train_ds, self.scorer)
 
         simple_score = scorer(_DummyModel, Dataset(pd.DataFrame(simple_pred), label=y_test))
-        pred_score = scorer(model, Dataset(test_ds.features_columns, label=y_test))
+        pred_score = scorer(model, Dataset(test_ds.features_columns, label=y_test, cat_features=test_ds.cat_features))
 
         return simple_score, pred_score, scorer.name