Export MulticlassLabelDataset from aif360.datasets

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

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

docs/source/modules/datasets.rst

@@ -21,6 +21,7 @@ Base classes
    datasets.Dataset
    datasets.StructuredDataset
    datasets.BinaryLabelDataset
+   datasets.MulticlassLabelDataset
    datasets.StandardDataset
    datasets.RegressionDataset
 

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]