58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
from cleanlab.experimental.span_classification import (
|
|
find_label_issues,
|
|
display_issues,
|
|
get_label_quality_scores,
|
|
)
|
|
import numpy as np
|
|
import pytest
|
|
import warnings
|
|
|
|
warnings.filterwarnings("ignore")
|
|
words = [["I", "love", "Cleanlab", "Inc"], ["A", "new", "park"]]
|
|
|
|
pred_probs = [
|
|
np.array([0.3, 0.2, 0.9, 0.1]),
|
|
np.array([0.1, 0.1, 0.9]),
|
|
]
|
|
|
|
labels = [[0, 0, 1, 1], [0, 0, 1]]
|
|
class_names = ["O", "Span"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_labels",
|
|
[labels, [np.array(l) for l in labels]],
|
|
ids=["list labels", "np.array labels"],
|
|
)
|
|
@pytest.mark.filterwarnings("ignore::DeprecationWarning")
|
|
def test_find_label_issues(test_labels):
|
|
issues = find_label_issues(test_labels, pred_probs)
|
|
assert isinstance(issues, list)
|
|
assert len(issues) == 1
|
|
assert issues[0] == (0, 3)
|
|
|
|
|
|
issues = find_label_issues(labels, pred_probs)
|
|
|
|
|
|
def test_display_issues():
|
|
display_issues(issues, words)
|
|
display_issues(issues, tokens=words, labels=labels)
|
|
display_issues(issues, words, pred_probs=pred_probs)
|
|
display_issues(issues, words, pred_probs=pred_probs, labels=labels)
|
|
display_issues(issues, words, pred_probs=pred_probs, labels=labels, class_names=class_names)
|
|
|
|
|
|
@pytest.fixture(name="label_quality_scores")
|
|
def fixture_label_quality_scores():
|
|
sentence_scores, token_info = get_label_quality_scores(labels, pred_probs)
|
|
return sentence_scores, token_info
|
|
|
|
|
|
def test_get_label_quality_scores(label_quality_scores):
|
|
sentence_scores, token_info = label_quality_scores
|
|
assert len(sentence_scores) == 2
|
|
assert np.allclose(sentence_scores, [0.1, 0.9])
|
|
assert len(token_info) == 2
|
|
assert np.allclose(token_info[0], [0.7, 0.8, 0.9, 0.1])
|