Files
cleanlab--cleanlab/tests/test_span_classification.py
T
2026-07-13 12:49:22 +08:00

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])