mirror of
https://github.com/IBM/ai-privacy-toolkit.git
synced 2026-05-07 02:52:39 +02:00
45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
|
import pytest
|
||
|
|
|
||
|
|
from apt.utils.models import SklearnClassifier, SklearnRegressor
|
||
|
|
from apt.utils import dataset_utils
|
||
|
|
|
||
|
|
from sklearn.tree import DecisionTreeRegressor
|
||
|
|
from sklearn.ensemble import RandomForestClassifier
|
||
|
|
|
||
|
|
def test_sklearn_classifier():
|
||
|
|
(x_train, y_train), (x_test, y_test) = dataset_utils.get_iris_dataset()
|
||
|
|
underlying_model = RandomForestClassifier()
|
||
|
|
model = SklearnClassifier(underlying_model)
|
||
|
|
model.fit(x_train, y_train)
|
||
|
|
pred = model.predict(x_test)
|
||
|
|
assert(pred.shape[0] == x_test.shape[0])
|
||
|
|
|
||
|
|
score = model.score(x_test, y_test)
|
||
|
|
assert(0.0 <= score <= 1.0)
|
||
|
|
|
||
|
|
def test_sklearn_regressor():
|
||
|
|
(x_train, y_train), (x_test, y_test) = dataset_utils.get_diabetes_dataset()
|
||
|
|
underlying_model = DecisionTreeRegressor()
|
||
|
|
model = SklearnRegressor(underlying_model)
|
||
|
|
model.fit(x_train, y_train)
|
||
|
|
pred = model.predict(x_test)
|
||
|
|
assert (pred.shape[0] == x_test.shape[0])
|
||
|
|
|
||
|
|
score = model.score(x_test, y_test)
|
||
|
|
|
||
|
|
losses = model.loss(x_test, y_test)
|
||
|
|
assert (losses.shape[0] == x_test.shape[0])
|
||
|
|
|
||
|
|
|
||
|
|
# Probably not needed for now, as we will not be using these wrappers directly in ART.
|
||
|
|
# def test_sklearn_decision_tree():
|
||
|
|
# (x_train, y_train), (x_test, y_test) = dataset_utils.get_iris_dataset()
|
||
|
|
# underlying_model = DecisionTreeClassifier()
|
||
|
|
# model = SklearnDecisionTreeClassifier(underlying_model)
|
||
|
|
# model.fit(x_train, y_train)
|
||
|
|
# pred = model.predict(x_test)
|
||
|
|
# assert(pred.shape[0] == x_test.shape[0])
|
||
|
|
#
|
||
|
|
# score = model.score(x_test, y_test)
|
||
|
|
# assert(0.0 <= score <= 1.0)
|