diff --git a/aif360/datasets/__init__.py b/aif360/datasets/__init__.py index a9d17caf..fe4f0a32 100644 --- a/aif360/datasets/__init__.py +++ b/aif360/datasets/__init__.py @@ -1,6 +1,7 @@ from aif360.datasets.dataset import Dataset from aif360.datasets.structured_dataset import StructuredDataset from aif360.datasets.binary_label_dataset import BinaryLabelDataset +from aif360.datasets.multiclass_label_dataset import MulticlassLabelDataset from aif360.datasets.standard_dataset import StandardDataset from aif360.datasets.adult_dataset import AdultDataset from aif360.datasets.bank_dataset import BankDataset diff --git a/aif360/datasets/multiclass_label_dataset.py b/aif360/datasets/multiclass_label_dataset.py index 141e3dcf..e8245d58 100644 --- a/aif360/datasets/multiclass_label_dataset.py +++ b/aif360/datasets/multiclass_label_dataset.py @@ -24,7 +24,7 @@ def __init__(self, favorable_label = [1.], unfavorable_label = [0.], **kwargs): def validate_dataset(self): """Error checking and type validation. -​ + Raises: ValueError: `labels` must be shape [n, 1]. ValueError: `favorable_label` and `unfavorable_label` must be the diff --git a/docs/source/modules/datasets.rst b/docs/source/modules/datasets.rst index 5c11921a..b89bd387 100644 --- a/docs/source/modules/datasets.rst +++ b/docs/source/modules/datasets.rst @@ -21,6 +21,7 @@ Base classes datasets.Dataset datasets.StructuredDataset datasets.BinaryLabelDataset + datasets.MulticlassLabelDataset datasets.StandardDataset datasets.RegressionDataset diff --git a/tests/test_multiclass_label_dataset.py b/tests/test_multiclass_label_dataset.py new file mode 100644 index 00000000..c3b4adb8 --- /dev/null +++ b/tests/test_multiclass_label_dataset.py @@ -0,0 +1,21 @@ +import numpy as np +import pandas as pd + +from aif360.datasets import MulticlassLabelDataset +from aif360.datasets.multiclass_label_dataset import ( + MulticlassLabelDataset as MulticlassLabelDatasetFromModule) + + +def test_importable_from_datasets_package(): + assert MulticlassLabelDataset is MulticlassLabelDatasetFromModule + + +def test_init_from_package_import(): + df = pd.DataFrame({'feat': [0, 1, 2, 1], 'label': [0, 1, 2, 2]}) + mcld = MulticlassLabelDataset(favorable_label=[0, 1], unfavorable_label=[2], + df=df, label_names=['label'], + protected_attribute_names=['feat']) + + assert np.all(mcld.labels == df[['label']].values) + assert mcld.favorable_label == [0, 1] + assert mcld.unfavorable_label == [2]