53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
import pytest
|
|
from cleanlab.datalab.internal.data import Data
|
|
from cleanlab.datalab.internal.data_issues import DataIssues, _ClassificationInfoStrategy
|
|
from cleanlab.datalab.internal.task import Task
|
|
|
|
|
|
class TestDataIssues:
|
|
labels = ["B", "A", "B"]
|
|
label_name = "labels"
|
|
strategy = _ClassificationInfoStrategy
|
|
|
|
@pytest.fixture
|
|
def data_issues(self):
|
|
data = Data(
|
|
data={self.label_name: self.labels},
|
|
task=Task.CLASSIFICATION,
|
|
label_name=self.label_name,
|
|
)
|
|
data_issues = DataIssues(data=data, strategy=self.strategy)
|
|
yield data_issues
|
|
|
|
def test_data_issues_init(self, data_issues):
|
|
assert hasattr(data_issues, "issues")
|
|
assert hasattr(data_issues, "issue_summary")
|
|
assert hasattr(data_issues, "info")
|
|
|
|
def test_statistics(self, data_issues):
|
|
stats = data_issues.statistics
|
|
|
|
assert stats == data_issues.info["statistics"]
|
|
assert stats["num_examples"] == 3, f"Incorrect number of examples: {stats['num_examples']}"
|
|
assert stats["class_names"] == ["A", "B"], f"Incorrect class names: {stats['class_names']}"
|
|
assert stats["num_classes"] == 2, f"Incorrect number of classes: {stats['num_classes']}"
|
|
assert stats["multi_label"] is False
|
|
assert (
|
|
stats["health_score"] is None
|
|
), f"Health score should initially be None, but is {stats['health_score']}"
|
|
|
|
def test_get_info(self, data_issues):
|
|
with pytest.raises(ValueError):
|
|
data_issues.get_info("nonexistent_issue")
|
|
assert data_issues.get_info("statistics") == data_issues.info["statistics"]
|
|
|
|
def test_get_info_label(self, data_issues):
|
|
data_issues.info["label"] = {"given_label": [0, 1, 1], "predicted_label": [1, 0, 1]}
|
|
info = data_issues.get_info("label")
|
|
|
|
label_format_error_message = (
|
|
"get_info('label') should return the given label formatted with the class names"
|
|
)
|
|
assert info.get("given_label").tolist() == ["A", "B", "B"], label_format_error_message
|
|
assert info.get("predicted_label").tolist() == self.labels, label_format_error_message
|