mirror of
https://github.com/IBM/ai-privacy-toolkit.git
synced 2026-05-10 04:22:37 +02:00
fix bug and update test_model
This commit is contained in:
parent
fd9f134336
commit
300e391432
3 changed files with 19 additions and 13 deletions
|
|
@ -66,6 +66,7 @@ class Anonymize:
|
|||
self.anonymizer = DecisionTreeRegressor(random_state=10, min_samples_split=2, min_samples_leaf=self.k)
|
||||
else:
|
||||
self.anonymizer = DecisionTreeClassifier(random_state=10, min_samples_split=2, min_samples_leaf=self.k)
|
||||
|
||||
self.anonymizer.fit(x_prepared, y)
|
||||
cells_by_id = self._calculate_cells(x, x_prepared)
|
||||
return self._anonymize_data_numpy(x, x_prepared, cells_by_id)
|
||||
|
|
@ -80,6 +81,8 @@ class Anonymize:
|
|||
self.anonymizer = DecisionTreeRegressor(random_state=10, min_samples_split=2, min_samples_leaf=self.k)
|
||||
else:
|
||||
self.anonymizer = DecisionTreeClassifier(random_state=10, min_samples_split=2, min_samples_leaf=self.k)
|
||||
if len(y.shape) > 1:
|
||||
y = np.argmax(y, axis=1)
|
||||
self.anonymizer.fit(x_prepared, y)
|
||||
cells_by_id = self._calculate_cells(x, x_prepared)
|
||||
return self._anonymize_data_pandas(x, x_prepared, cells_by_id)
|
||||
|
|
|
|||
|
|
@ -46,7 +46,10 @@ class SklearnClassifier(SklearnModel):
|
|||
:type y: `np.ndarray` or `pandas.DataFrame`
|
||||
"""
|
||||
encoder = OneHotEncoder(sparse=False)
|
||||
y_encoded = encoder.fit_transform(y.reshape(-1, 1))
|
||||
if type(y) == np.ndarray:
|
||||
y_encoded = encoder.fit_transform(y.reshape(-1, 1))
|
||||
else:
|
||||
y_encoded = encoder.fit_transform(y.values.reshape(-1, 1))
|
||||
self._art_model.fit(x, y_encoded, **kwargs)
|
||||
|
||||
def predict(self, x: np.ndarray, **kwargs) -> np.ndarray:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue