diff --git a/apt/utils/dataset_utils.py b/apt/utils/dataset_utils.py index cf1a701..f99c6cc 100644 --- a/apt/utils/dataset_utils.py +++ b/apt/utils/dataset_utils.py @@ -5,8 +5,6 @@ import ssl from os import path, mkdir from six.moves.urllib.request import urlretrieve -from apt.utils.datasets import ArrayDataset, Data - def _load_iris(test_set_size: float = 0.3): iris = datasets.load_iris() @@ -16,10 +14,8 @@ def _load_iris(test_set_size: float = 0.3): # Split training and test sets x_train, x_test, y_train, y_test = model_selection.train_test_split(data, labels, test_size=test_set_size, random_state=18, stratify=labels) - train_dataset = ArrayDataset(x_train, y_train) - test_dataset = ArrayDataset(x_test, y_test) - dataset = Data(train_dataset, test_dataset) - return dataset + + return (x_train, y_train), (x_test, y_test) def get_iris_dataset(): @@ -41,10 +37,7 @@ def _load_diabetes(test_set_size: float = 0.3): x_train, x_test, y_train, y_test = model_selection.train_test_split(data, labels, test_size=test_set_size, random_state=18) - train_dataset = ArrayDataset(x_train, y_train) - test_dataset = ArrayDataset(x_test, y_test) - dataset = Data(train_dataset, test_dataset) - return dataset + return (x_train, y_train), (x_test, y_test) def get_diabetes_dataset(): @@ -104,10 +97,7 @@ def get_german_credit_dataset(test_set: float = 0.3): x_test.reset_index(drop=True, inplace=True) y_test.reset_index(drop=True, inplace=True) - train_dataset = ArrayDataset(x_train, y_train) - test_dataset = ArrayDataset(x_test, y_test) - dataset = Data(train_dataset, test_dataset) - return dataset + return (x_train, y_train), (x_test, y_test) def _modify_german_dataset(data): @@ -166,10 +156,8 @@ def get_adult_dataset(): y_train = train.loc[:, 'label'] x_test = test.drop(['label'], axis=1) y_test = test.loc[:, 'label'] - train_dataset = ArrayDataset(x_train, y_train) - test_dataset = ArrayDataset(x_test, y_test) - dataset = Data(train_dataset, test_dataset) - return dataset + + return (x_train, y_train), (x_test, y_test) def _modify_adult_dataset(data): @@ -327,10 +315,5 @@ def get_nursery_dataset(raw: bool = True, test_set: float = 0.2, transform_socia y_train = train.loc[:, "label"] x_test = test.drop(["label"], axis=1) y_test = test.loc[:, "label"] - x_train = x_train.astype(str) - x_test = x_test.astype(str) - train_dataset = ArrayDataset(x_train, y_train) - test_dataset = ArrayDataset(x_test, y_test) - dataset = Data(train_dataset, test_dataset) - return dataset + return (x_train, y_train), (x_test, y_test) diff --git a/tests/test_anonymizer.py b/tests/test_anonymizer.py index da38655..4a96b9a 100644 --- a/tests/test_anonymizer.py +++ b/tests/test_anonymizer.py @@ -7,29 +7,29 @@ from apt.anonymization import Anonymize from apt.utils.dataset_utils import get_iris_dataset, get_adult_dataset, get_nursery_dataset from sklearn.datasets import load_diabetes from sklearn.model_selection import train_test_split -from apt.utils.datasets import ArrayDataset, Data + def test_anonymize_ndarray_iris(): - dataset = get_iris_dataset() + (x_train, y_train), _ = get_iris_dataset() model = DecisionTreeClassifier() - model.fit(dataset.get_train_samples(), dataset.get_train_labels()) - pred = model.predict(dataset.get_train_samples()) + model.fit(x_train, y_train) + pred = model.predict(x_train) k = 10 QI = [0, 2] anonymizer = Anonymize(k, QI) - anon = anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), pred)) - assert(len(np.unique(anon[:, QI], axis=0)) < len(np.unique(dataset.get_train_samples()[:, QI], axis=0))) + anon = anonymizer.anonymize(x_train, pred) + assert(len(np.unique(anon[:, QI], axis=0)) < len(np.unique(x_train[:, QI], axis=0))) _, counts_elements = np.unique(anon[:, QI], return_counts=True) assert (np.min(counts_elements) >= k) - assert ((np.delete(anon, QI, axis=1) == np.delete(dataset.get_train_samples(), QI, axis=1)).all()) + assert ((np.delete(anon, QI, axis=1) == np.delete(x_train, QI, axis=1)).all()) def test_anonymize_pandas_adult(): - dataset = get_adult_dataset() - encoded = OneHotEncoder().fit_transform(dataset.get_train_samples()) + (x_train, y_train), _ = get_adult_dataset() + encoded = OneHotEncoder().fit_transform(x_train) model = DecisionTreeClassifier() - model.fit(encoded, dataset.get_train_labels()) + model.fit(encoded, y_train) pred = model.predict(encoded) k = 100 @@ -38,51 +38,51 @@ def test_anonymize_pandas_adult(): categorical_features = ['workclass', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'native-country'] anonymizer = Anonymize(k, QI, categorical_features=categorical_features) - anon = anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), pred)) + anon = anonymizer.anonymize(x_train, pred) - assert(anon.loc[:, QI].drop_duplicates().shape[0] < dataset.get_train_samples().loc[:, QI].drop_duplicates().shape[0]) + assert(anon.loc[:, QI].drop_duplicates().shape[0] < x_train.loc[:, QI].drop_duplicates().shape[0]) assert (anon.loc[:, QI].value_counts().min() >= k) - assert (anon.drop(QI, axis=1).equals(dataset.get_train_samples().drop(QI, axis=1))) + assert (anon.drop(QI, axis=1).equals(x_train.drop(QI, axis=1))) def test_anonymize_pandas_nursery(): - dataset = get_nursery_dataset() - encoded = OneHotEncoder().fit_transform(dataset.get_train_samples()) + (x_train, y_train), _ = get_nursery_dataset() + x_train = x_train.astype(str) + encoded = OneHotEncoder().fit_transform(x_train) model = DecisionTreeClassifier() - model.fit(encoded, dataset.get_train_labels()) + model.fit(encoded, y_train) pred = model.predict(encoded) k = 100 QI = ["finance", "social", "health"] categorical_features = ["parents", "has_nurs", "form", "housing", "finance", "social", "health", 'children'] anonymizer = Anonymize(k, QI, categorical_features=categorical_features) - anon = anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), pred)) + anon = anonymizer.anonymize(x_train, pred) - assert(anon.loc[:, QI].drop_duplicates().shape[0] < dataset.get_train_samples().loc[:, QI].drop_duplicates().shape[0]) + assert(anon.loc[:, QI].drop_duplicates().shape[0] < x_train.loc[:, QI].drop_duplicates().shape[0]) assert (anon.loc[:, QI].value_counts().min() >= k) - assert (anon.drop(QI, axis=1).equals(dataset.get_train_samples().drop(QI, axis=1))) + assert (anon.drop(QI, axis=1).equals(x_train.drop(QI, axis=1))) def test_regression(): - x_train, x_test, y_train, y_test = train_test_split(load_diabetes().data, load_diabetes().target, test_size=0.5, random_state=14) - train_dataset = ArrayDataset(x_train, y_train) - test_dataset = ArrayDataset(x_test, y_test) - dataset = Data(train_dataset, test_dataset) + dataset = load_diabetes() + x_train, x_test, y_train, y_test = train_test_split(dataset.data, dataset.target, test_size=0.5, random_state=14) + model = DecisionTreeRegressor(random_state=10, min_samples_split=2) - model.fit(dataset.get_train_samples(), dataset.get_train_labels()) - pred = model.predict(dataset.get_train_samples()) + model.fit(x_train, y_train) + pred = model.predict(x_train) k = 10 QI = [0, 2, 5, 8] anonymizer = Anonymize(k, QI, is_regression=True) - anon = anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), pred)) - print('Base model accuracy (R2 score): ', model.score(dataset.get_test_samples(), dataset.get_test_labels())) - model.fit(anon, dataset.get_train_labels()) - print('Base model accuracy (R2 score) after anonymization: ', model.score(dataset.get_test_samples(), dataset.get_test_labels())) - assert(len(np.unique(anon[:, QI], axis=0)) < len(np.unique(dataset.get_train_samples()[:, QI], axis=0))) + anon = anonymizer.anonymize(x_train, pred) + print('Base model accuracy (R2 score): ', model.score(x_test, y_test)) + model.fit(anon, y_train) + print('Base model accuracy (R2 score) after anonymization: ', model.score(x_test, y_test)) + assert(len(np.unique(anon[:, QI], axis=0)) < len(np.unique(x_train[:, QI], axis=0))) _, counts_elements = np.unique(anon[:, QI], return_counts=True) assert (np.min(counts_elements) >= k) - assert ((np.delete(anon, QI, axis=1) == np.delete(dataset.get_train_samples(), QI, axis=1)).all()) + assert ((np.delete(anon, QI, axis=1) == np.delete(x_train, QI, axis=1)).all()) def test_errors(): @@ -93,9 +93,9 @@ def test_errors(): with pytest.raises(ValueError): Anonymize(2, None) anonymizer = Anonymize(10, [0, 2]) - dataset = get_iris_dataset() + (x_train, y_train), (x_test, y_test) = get_iris_dataset() with pytest.raises(ValueError): - anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), dataset.get_test_labels())) - dataset = get_adult_dataset() + anonymizer.anonymize(x_train, y_test) + (x_train, y_train), _ = get_adult_dataset() with pytest.raises(ValueError): - anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), dataset.get_train_labels())) + anonymizer.anonymize(x_train, y_train)