diff --git a/apt/anonymization/anonymizer.py b/apt/anonymization/anonymizer.py index d7452b1..23599b6 100644 --- a/apt/anonymization/anonymizer.py +++ b/apt/anonymization/anonymizer.py @@ -5,7 +5,7 @@ from collections import Counter from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor from sklearn.preprocessing import OneHotEncoder -from apt.utils.datasets import BaseDataset, Data +from apt.utils.datasets import ArrayDataset from typing import Union, Optional @@ -38,7 +38,7 @@ class Anonymize: self.categorical_features = categorical_features self.is_regression = is_regression - def anonymize(self, dataset: BaseDataset) -> Union[np.ndarray, pd.DataFrame]: + def anonymize(self, dataset: ArrayDataset) -> Union[np.ndarray, pd.DataFrame]: """ Method for performing model-guided anonymization. diff --git a/apt/utils/dataset_utils.py b/apt/utils/dataset_utils.py index 3bafa7b..cf1a701 100644 --- a/apt/utils/dataset_utils.py +++ b/apt/utils/dataset_utils.py @@ -5,7 +5,7 @@ import ssl from os import path, mkdir from six.moves.urllib.request import urlretrieve -from apt.utils.datasets import BaseDataset, Data +from apt.utils.datasets import ArrayDataset, Data def _load_iris(test_set_size: float = 0.3): @@ -16,8 +16,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 = BaseDataset(x_train, y_train) - test_dataset = BaseDataset(x_test, y_test) + train_dataset = ArrayDataset(x_train, y_train) + test_dataset = ArrayDataset(x_test, y_test) dataset = Data(train_dataset, test_dataset) return dataset @@ -41,8 +41,8 @@ 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 = BaseDataset(x_train, y_train) - test_dataset = BaseDataset(x_test, y_test) + train_dataset = ArrayDataset(x_train, y_train) + test_dataset = ArrayDataset(x_test, y_test) dataset = Data(train_dataset, test_dataset) return dataset @@ -104,8 +104,8 @@ 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 = BaseDataset(x_train, y_train) - test_dataset = BaseDataset(x_test, y_test) + train_dataset = ArrayDataset(x_train, y_train) + test_dataset = ArrayDataset(x_test, y_test) dataset = Data(train_dataset, test_dataset) return dataset @@ -166,8 +166,8 @@ def get_adult_dataset(): y_train = train.loc[:, 'label'] x_test = test.drop(['label'], axis=1) y_test = test.loc[:, 'label'] - train_dataset = BaseDataset(x_train, y_train) - test_dataset = BaseDataset(x_test, y_test) + train_dataset = ArrayDataset(x_train, y_train) + test_dataset = ArrayDataset(x_test, y_test) dataset = Data(train_dataset, test_dataset) return dataset @@ -330,7 +330,7 @@ def get_nursery_dataset(raw: bool = True, test_set: float = 0.2, transform_socia x_train = x_train.astype(str) x_test = x_test.astype(str) - train_dataset = BaseDataset(x_train, y_train) - test_dataset = BaseDataset(x_test, y_test) + train_dataset = ArrayDataset(x_train, y_train) + test_dataset = ArrayDataset(x_test, y_test) dataset = Data(train_dataset, test_dataset) return dataset diff --git a/tests/test_anonymizer.py b/tests/test_anonymizer.py index 280309e..da38655 100644 --- a/tests/test_anonymizer.py +++ b/tests/test_anonymizer.py @@ -7,7 +7,7 @@ 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 BaseDataset, Data +from apt.utils.datasets import ArrayDataset, Data def test_anonymize_ndarray_iris(): dataset = get_iris_dataset() @@ -18,7 +18,7 @@ def test_anonymize_ndarray_iris(): k = 10 QI = [0, 2] anonymizer = Anonymize(k, QI) - anon = anonymizer.anonymize(BaseDataset(dataset.get_train_samples(), pred)) + 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))) _, counts_elements = np.unique(anon[:, QI], return_counts=True) assert (np.min(counts_elements) >= k) @@ -38,7 +38,7 @@ 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(BaseDataset(dataset.get_train_samples(), pred)) + anon = anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), pred)) assert(anon.loc[:, QI].drop_duplicates().shape[0] < dataset.get_train_samples().loc[:, QI].drop_duplicates().shape[0]) assert (anon.loc[:, QI].value_counts().min() >= k) @@ -56,7 +56,7 @@ def test_anonymize_pandas_nursery(): 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(BaseDataset(dataset.get_train_samples(), pred)) + anon = anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), pred)) assert(anon.loc[:, QI].drop_duplicates().shape[0] < dataset.get_train_samples().loc[:, QI].drop_duplicates().shape[0]) assert (anon.loc[:, QI].value_counts().min() >= k) @@ -66,8 +66,8 @@ def test_anonymize_pandas_nursery(): 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 = BaseDataset(x_train, y_train) - test_dataset = BaseDataset(x_test, y_test) + train_dataset = ArrayDataset(x_train, y_train) + test_dataset = ArrayDataset(x_test, y_test) dataset = Data(train_dataset, test_dataset) model = DecisionTreeRegressor(random_state=10, min_samples_split=2) model.fit(dataset.get_train_samples(), dataset.get_train_labels()) @@ -75,7 +75,7 @@ def test_regression(): k = 10 QI = [0, 2, 5, 8] anonymizer = Anonymize(k, QI, is_regression=True) - anon = anonymizer.anonymize(BaseDataset(dataset.get_train_samples(), pred)) + 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())) @@ -95,7 +95,7 @@ def test_errors(): anonymizer = Anonymize(10, [0, 2]) dataset = get_iris_dataset() with pytest.raises(ValueError): - anonymizer.anonymize(BaseDataset(dataset.get_train_samples(), dataset.get_test_labels())) + anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), dataset.get_test_labels())) dataset = get_adult_dataset() with pytest.raises(ValueError): - anonymizer.anonymize(BaseDataset(dataset.get_train_samples(), dataset.get_train_labels())) + anonymizer.anonymize(ArrayDataset(dataset.get_train_samples(), dataset.get_train_labels())) diff --git a/tests/test_model.py b/tests/test_model.py index 7dd151c..1ef5b34 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -9,28 +9,24 @@ from sklearn.ensemble import RandomForestClassifier def test_sklearn_classifier(): - (x_train, y_train), (x_test, y_test) = dataset_utils.get_iris_dataset() + dataset = dataset_utils.get_iris_dataset() underlying_model = RandomForestClassifier() model = SklearnClassifier(underlying_model) - train = ArrayDataset(x_train, y_train) - test = ArrayDataset(x_test, y_test) - model.fit(train) - pred = model.predict(x_test) - assert(pred.shape[0] == x_test.shape[0]) + model.fit(dataset.train) + pred = model.predict(dataset.get_test_samples()) + assert(pred.shape[0] == dataset.get_test_samples().shape[0]) - score = model.score(test) + score = model.score(dataset.test) assert(0.0 <= score <= 1.0) def test_sklearn_regressor(): - (x_train, y_train), (x_test, y_test) = dataset_utils.get_diabetes_dataset() + dataset = dataset_utils.get_diabetes_dataset() underlying_model = DecisionTreeRegressor() model = SklearnRegressor(underlying_model) - train = ArrayDataset(x_train, y_train) - test = ArrayDataset(x_test, y_test) - model.fit(train) - pred = model.predict(x_test) - assert (pred.shape[0] == x_test.shape[0]) + model.fit(dataset.train) + pred = model.predict(dataset.get_test_samples()) + assert (pred.shape[0] == dataset.get_test_samples().shape[0]) - score = model.score(test) + score = model.score(dataset.test) assert (0 <= score <= 1)