49 lines
2.1 KiB
Python
49 lines
2.1 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from cleanlab import Datalab
|
|
from cleanlab.datalab.internal.issue_manager.multilabel.label import MultilabelIssueManager
|
|
from cleanlab.internal.multilabel_utils import onehot2int
|
|
|
|
|
|
class TestLabelIssueManager:
|
|
@pytest.fixture
|
|
def data(self):
|
|
# True labels of multilabel dataset
|
|
np.random.seed(10)
|
|
true_y = np.array(
|
|
[[0, 1, 0], [0, 1, 1], [1, 1, 1], [1, 0, 1], [0, 1, 1], [1, 1, 0], [1, 1, 0]]
|
|
)
|
|
pred_probs = np.full_like(true_y, fill_value=0.1, dtype=float)
|
|
pred_probs[true_y == 1] = 0.9
|
|
# Flip labels of some rows to add noise
|
|
candidate_rows = np.where(true_y.sum(axis=1) < 3)[0]
|
|
noisy_rows = np.random.choice(candidate_rows, 2, replace=False)
|
|
noisy_y = true_y.copy()
|
|
noisy_y[noisy_rows] = 1 - true_y[noisy_rows]
|
|
labels = onehot2int(noisy_y)
|
|
return {"labels": labels, "pred_probs": pred_probs}
|
|
|
|
@pytest.fixture
|
|
def issue_manager(self, data):
|
|
labels = data["labels"]
|
|
lab = Datalab({"labels": labels}, task="multilabel", label_name="labels")
|
|
return MultilabelIssueManager(datalab=lab)
|
|
|
|
def test_find_issues(self, data, issue_manager):
|
|
"""Test that the find_issues method works."""
|
|
pred_probs = data["pred_probs"]
|
|
issue_manager.find_issues(pred_probs=pred_probs)
|
|
issues, summary, info = issue_manager.issues, issue_manager.summary, issue_manager.info
|
|
assert isinstance(issues, pd.DataFrame), "Issues should be a dataframe"
|
|
assert isinstance(summary, pd.DataFrame), "Summary should be a dataframe"
|
|
assert summary["issue_type"].values[0] == "label"
|
|
assert issues.index[issues["is_label_issue"]].tolist() == [3, 6]
|
|
assert pytest.approx(summary["score"].values[0], abs=1e-3) == 0.6714
|
|
assert isinstance(info, dict), "Info should be a dict"
|
|
|
|
issue_manager.find_issues(pred_probs=pred_probs, frac_noise=0.5)
|
|
issues = issue_manager.issues
|
|
assert issues.index[issues["is_label_issue"]].tolist() == [3]
|