wehub-resource-sync
96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
import numpy as np
|
|
import pytest
|
|
from hypothesis import given, settings, strategies as st
|
|
from hypothesis.strategies import composite
|
|
from hypothesis.extra.numpy import arrays
|
|
|
|
from sklearn.neighbors import NearestNeighbors
|
|
|
|
from cleanlab.data_valuation import _knn_shapley_score, data_shapley_knn
|
|
from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
|
|
|
|
|
|
class TestDataValuation:
|
|
K = 3
|
|
N = 100
|
|
num_features = 10
|
|
|
|
@pytest.fixture
|
|
def features(self):
|
|
return np.random.rand(self.N, self.num_features)
|
|
|
|
@pytest.fixture
|
|
def labels(self):
|
|
return np.random.randint(0, 2, self.N)
|
|
|
|
@pytest.fixture
|
|
def knn_graph(self, features):
|
|
knn = NearestNeighbors(n_neighbors=self.K).fit(features)
|
|
knn_graph = knn.kneighbors_graph(mode="distance")
|
|
return knn_graph
|
|
|
|
def test_data_shapley_knn(self, labels, features):
|
|
shapley = data_shapley_knn(labels, features=features, k=self.K)
|
|
assert shapley.shape == (100,)
|
|
assert np.all(shapley >= 0)
|
|
assert np.all(shapley <= 1)
|
|
|
|
def test_data_shapley_knn_with_knn_graph(self, labels, knn_graph):
|
|
shapley = data_shapley_knn(labels, knn_graph=knn_graph, k=self.K)
|
|
assert shapley.shape == (100,)
|
|
assert np.all(shapley >= 0)
|
|
assert np.all(shapley <= 1)
|
|
|
|
|
|
@composite
|
|
def valid_data(draw):
|
|
"""
|
|
A custom strategy to generate valid labels, features, and k such that:
|
|
- labels and features have the same length
|
|
- k is less than the length of labels and features
|
|
"""
|
|
# Generate a valid length for labels and features
|
|
length = draw(st.integers(min_value=11, max_value=1000))
|
|
|
|
# Generate labels and features of the same length
|
|
labels = draw(
|
|
arrays(
|
|
dtype=np.int32,
|
|
shape=length,
|
|
elements=st.integers(min_value=0, max_value=length - 1),
|
|
)
|
|
)
|
|
features = draw(
|
|
arrays(
|
|
dtype=np.float64,
|
|
shape=(length, draw(st.integers(min_value=2, max_value=50))),
|
|
elements=st.floats(min_value=-1.0, max_value=1.0),
|
|
)
|
|
)
|
|
|
|
# Generate k such that it is less than the length of labels and features
|
|
k = draw(st.integers(min_value=1, max_value=length - 1))
|
|
|
|
return labels, features, k
|
|
|
|
|
|
class TestDataShapleyKNNScore:
|
|
"""This test class prioritizes testing the raw/untransformed outputs of the _knn_shapley_score function."""
|
|
|
|
@settings(
|
|
max_examples=1000, deadline=None
|
|
) # Increase the number of examples to test more cases
|
|
@given(valid_data())
|
|
def test_knn_shapley_score_property(self, data):
|
|
labels, features, k = data
|
|
|
|
knn_graph, _ = create_knn_graph_and_index(features, n_neighbors=k)
|
|
neighbor_indices = knn_graph.indices.reshape(-1, k)
|
|
|
|
scores = _knn_shapley_score(neighbor_indices, labels, k)
|
|
|
|
# Shapley scores should be between -1 and 1
|
|
assert scores.shape == (len(labels),)
|
|
assert np.all(scores >= -1)
|
|
assert np.all(scores <= 1)
|