Modify headers when running in suite on SingleDataset checks (deepche…

…cks#53) Co-authored-by: Matan Perlmutter <[email protected]>
harveyyeung · Oct 19, 2021 · f862dd6 · f862dd6
1 parent 43845fc
commit f862dd6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/mlchecks/base/suite.py b/mlchecks/base/suite.py
@@ -60,9 +60,13 @@ def run(self, model=None, train_dataset=None, validation_dataset=None, check_dat
                 results.append(check.run(dataset=train_dataset, compared_dataset=validation_dataset, model=model))
             elif isinstance(check, SingleDatasetBaseCheck):
                 if check_datasets_policy in ['both', 'train'] and train_dataset is not None:
-                    results.append(check.run(dataset=train_dataset))
+                    res = check.run(dataset=train_dataset)
+                    res.header = f'{res.header} - Train Dataset'
+                    results.append(res)
                 if check_datasets_policy in ['both', 'validation'] and validation_dataset is not None:
-                    results.append(check.run(dataset=validation_dataset))
+                    res = check.run(dataset=validation_dataset)
+                    res.header = f'{res.header} - Validation Dataset'
+                    results.append(res)
             elif isinstance(check, ModelOnlyBaseCheck):
                 results.append(check.run(model=model))
             elif isinstance(check, CheckSuite):

diff --git a/mlchecks/checks/integrity/__init__.py b/mlchecks/checks/integrity/__init__.py
@@ -2,4 +2,5 @@
 from .mixed_nulls import *
 from .string_mismatch import *
 from .mixed_types import *
+from .is_single_value import *
 from .invalid_chars import *