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

344 lines
12 KiB
Python

from cleanlab.internal.token_classification_utils import (
get_sentence,
filter_sentence,
process_token,
mapping,
merge_probs,
color_sentence,
_replace_sentence,
)
from cleanlab.token_classification.filter import find_label_issues
from cleanlab.token_classification.rank import (
get_label_quality_scores,
issues_from_scores,
_softmin_sentence_score,
)
from cleanlab.token_classification.summary import (
display_issues,
common_label_issues,
filter_by_token,
)
import numpy as np
import pandas as pd
import pytest
import warnings
warnings.filterwarnings("ignore")
words = [["Hello", "World"], ["#I", "love", "Cleanlab"], ["A"]]
sentences = ["Hello World", "#I love Cleanlab", "A"]
pred_probs = [
np.array([[0.9, 0.1, 0], [0.6, 0.2, 0.2]]),
np.array([[0.1, 0, 0.9], [0.1, 0.8, 0.1], [0.1, 0.8, 0.1]]),
np.array([[0.1, 0.1, 0.8]]),
]
labels = [[0, 0], [1, 1, 1], [2]]
maps = [0, 1, 0, 1]
class_names = ["A", "B", "C", "D"]
def test_get_sentence():
actual_sentences = list(map(get_sentence, words))
assert actual_sentences == sentences
# Test with allowed special characters
words_separated_by_hyphen = ["Heading", "-", "Title"]
assert get_sentence(words_separated_by_hyphen) == "Heading - Title"
words_within_parentheses = ["Some", "reason", "(", "Explanation", ")"]
assert get_sentence(words_within_parentheses) == "Some reason (Explanation)"
def test_filter_sentence():
filtered_sentences, mask = filter_sentence(sentences)
assert filtered_sentences == ["Hello World"]
assert mask == [True, False, False]
filtered_sentences, mask = filter_sentence(sentences, lambda x: len(x) > 1)
assert filtered_sentences == ["Hello World", "#I love Cleanlab"]
assert mask == [True, True, False]
filtered_sentences, mask = filter_sentence(sentences, lambda x: "#" not in x)
assert filtered_sentences == ["Hello World", "A"]
assert mask == [True, False, True]
def test_process_token():
test_cases = [
("Cleanlab", [("C", "a")], "aleanlab"),
("Cleanlab", [("C", "a"), ("a", "C")], "aleCnlCb"),
]
for token, replacements, expected in test_cases:
processed = process_token(token, replacements)
assert processed == expected
def test_mapping():
test_cases = [(l, expected) for l, expected in zip(labels, [[0, 0], [1, 1, 1], [0]])]
for l, expected in test_cases:
mapped = mapping(l, maps)
assert mapped == expected
def test_merge_probs():
merged_probs = merge_probs(pred_probs[0], maps)
expected = np.array([[0.9, 0.1], [0.8, 0.2]])
assert np.allclose(expected, merged_probs)
merged_probs = merge_probs(pred_probs[1], maps)
expected = np.array([[1.0, 0.0], [0.2, 0.8], [0.2, 0.8]])
assert np.allclose(expected, merged_probs)
merged_probs = merge_probs(pred_probs[2], maps)
expected = np.array([[0.9, 0.1]])
assert np.allclose(expected, merged_probs)
def test_merge_probs_with_normalization():
# Ignore probabilities for class/entity 0
norm_maps = [-1, 1, 0, 1]
merged_probs = merge_probs(pred_probs[0], norm_maps)
expected = np.array([[0.0, 1.0], [0.5, 0.5]])
assert np.allclose(expected, merged_probs)
merged_probs = merge_probs(pred_probs[1], norm_maps)
expected = np.array([[1.0, 0.0], [1 / 9, 8 / 9], [1 / 9, 8 / 9]])
assert np.allclose(expected, merged_probs)
merged_probs = merge_probs(pred_probs[2], norm_maps)
expected = np.array([[8 / 9, 1 / 9]])
# Ignore probabilities for class/entity 1
norm_maps = [0, -1, 0, 1]
merged_probs = merge_probs(pred_probs[0], norm_maps)
expected = np.array([[1.0, 0.0], [1.0, 0.0]])
assert np.allclose(expected, merged_probs)
merged_probs = merge_probs(pred_probs[1], norm_maps)
expected = np.array([[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]])
assert np.allclose(expected, merged_probs)
merged_probs = merge_probs(pred_probs[2], norm_maps)
expected = np.array([[1.0, 0.0]])
assert np.allclose(expected, merged_probs)
# Color boundaries
C_L, C_R = "\x1b[31m", "\x1b[0m"
@pytest.mark.parametrize(
"sentence,word,expected",
[
("Hello World", "World", f"Hello {C_L}World{C_R}"),
("Hello World", "help", "Hello World"),
("If you and I were to meet", "I", f"If you and {C_L}I{C_R} were to meet"),
("If you and I were to meet", "If you and I", f"{C_L}If you and I{C_R} were to meet"),
("If you and I were to meet", "If you and I w", f"{C_L}If you and I w{C_R}ere to meet"),
("I think I know this", "I", f"{C_L}I{C_R} think {C_L}I{C_R} know this"),
("A good reason for a test", "a", f"A good reason for {C_L}a{C_R} test"),
("ab ab a b ab", "ab a", f"ab {C_L}ab a{C_R} b ab"),
("ab ab ab ab", "ab a", f"{C_L}ab a{C_R}b {C_L}ab a{C_R}b"),
(
"Alan John Percivale (A.j.p.) Taylor died",
"(",
f"Alan John Percivale {C_L}({C_R}A.j.p.) Taylor died",
),
],
ids=[
"single_word",
"no_match",
"ignore_subwords",
"multi-token_match",
"substring_replacement",
"multiple_matches",
"case_sensitive",
"only_word_boundary",
"non_overlapping_substrings",
"issue_403-escape_special_regex_characters",
],
)
def test_color_sentence(monkeypatch: pytest.MonkeyPatch, sentence, word, expected):
import os
monkeypatch.setattr(os, "isatty", lambda fd: True)
monkeypatch.setattr("sys.stdout.isatty", lambda: True)
monkeypatch.setattr("sys.stdout.fileno", lambda: 1)
colored = color_sentence(sentence, word)
assert colored == expected
@pytest.mark.parametrize(
"sentence,word,expected",
[
("Hello World", "World", "Hello [EXPECTED]"),
("Hello World", "help", "Hello World"),
("If you and I were to meet", "I", "If you and [EXPECTED] were to meet"),
("If you and I were to meet", "If you and I", "[EXPECTED] were to meet"),
("If you and I were to meet", "If you and I w", "[EXPECTED]ere to meet"),
("I think I know this", "I", "[EXPECTED] think [EXPECTED] know this"),
("A good reason for a test", "a", "A good reason for [EXPECTED] test"),
("ab ab a b ab", "ab a", "ab [EXPECTED] b ab"),
("ab ab ab ab", "ab a", "[EXPECTED]b [EXPECTED]b"),
(
"Alan John Percivale (A.j.p.) Taylor died",
"(",
"Alan John Percivale [EXPECTED]A.j.p.) Taylor died",
),
],
ids=[
"single_word",
"no_match",
"ignore_subwords",
"multi-token_match",
"substring_replacement",
"multiple_matches",
"case_sensitive",
"only_word_boundary",
"non_overlapping_substrings",
"issue_403-escape_special_regex_characters",
],
)
def test_replace_sentence(sentence, word, expected):
new_sentence = _replace_sentence(sentence, word, "[EXPECTED]")
assert new_sentence == expected
issues = find_label_issues(labels, pred_probs)
@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] == (1, 0)
issues2 = find_label_issues(
test_labels, pred_probs, return_indices_ranked_by="normalized_margin", n_jobs=1
)
assert isinstance(issues2, list)
# Compare results with low_memory=True. Pass unused argument n_jobs=1
issues_lm = find_label_issues(test_labels, pred_probs, low_memory=True, n_jobs=1)
intersection = len(list(set(issues).intersection(set(issues_lm))))
union = len(set(issues)) + len(set(issues_lm)) - intersection
assert float(intersection) / union > 0.95
def test_softmin_sentence_score():
token_scores = [[0.9, 0.6], [0.0, 0.8, 0.8], [0.8]]
sentence_scores = _softmin_sentence_score(token_scores)
assert isinstance(sentence_scores, np.ndarray)
assert np.allclose(sentence_scores, [0.60074, 1.8e-07, 0.8])
# Temperature limits
sentence_scores = _softmin_sentence_score(token_scores, temperature=0)
assert np.allclose(sentence_scores, [0.6, 0.0, 0.8])
sentence_scores = _softmin_sentence_score(token_scores, temperature=np.inf)
assert np.allclose(sentence_scores, [0.75, 1.6 / 3, 0.8])
@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) == 3
assert np.allclose(sentence_scores, [0.6, 0, 0.8])
assert len(token_info) == 3
assert np.allclose(token_info[0], [0.9, 0.6])
sentence_scores_softmin, _ = get_label_quality_scores(
labels, pred_probs, sentence_score_method="softmin", tokens=words
)
assert len(sentence_scores_softmin) == 3
assert np.allclose(sentence_scores_softmin, [0.600741787, 1.8005624e-7, 0.8])
with pytest.raises(AssertionError) as excinfo:
get_label_quality_scores(
labels, pred_probs, sentence_score_method="unsupported_method", tokens=words
)
assert "Select from the following methods:" in str(excinfo.value)
def test_issues_from_scores(label_quality_scores):
sentence_scores, token_scores = label_quality_scores
issues = issues_from_scores(sentence_scores, token_scores=token_scores)
assert len(issues) == 1
assert issues[0] == (1, 0)
issues_without = issues_from_scores(sentence_scores)
assert len(issues_without) == 1
assert issues_without[0] == 1
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)
exclude = [(1, 2)] # Occurs in first token of second sentence "#I"
display_issues(issues, words, pred_probs=pred_probs, labels=labels, exclude=exclude)
top = 1
display_issues(issues, words, pred_probs=pred_probs, labels=labels, top=top)
issues_sentence_only = [i for i, _ in issues]
display_issues(issues_sentence_only, words)
TEST_KWARGS = {"labels": labels, "pred_probs": pred_probs, "class_names": class_names}
@pytest.mark.parametrize(
"test_issues",
[issues, issues + [(1, 0)]],
ids=["default issues", "augmented issues"],
)
@pytest.mark.parametrize(
"test_kwargs",
[
{},
TEST_KWARGS,
{**TEST_KWARGS, "top": 1},
{**TEST_KWARGS, "exclude": [(1, 2)]},
{**TEST_KWARGS, "verbose": False},
],
ids=["no kwargs", "labels+pred_probs+class_names", "...+top", "...+exclude", "...+no verbose"],
)
def test_common_label_issues(test_issues, test_kwargs):
df = common_label_issues(test_issues, words, **test_kwargs)
assert isinstance(df, pd.DataFrame)
columns = df.columns.tolist()
for col in ["token", "num_label_issues"]:
assert col in columns
if test_kwargs:
for col in ["given_label", "predicted_label"]:
assert col in columns
@pytest.mark.parametrize(
"test_token,expected_issues",
[
("Hello", []),
("#I", [(1, 0)]),
],
)
def test_filter_by_token(test_token, expected_issues):
returned_issues = filter_by_token(test_token, issues, words)
assert returned_issues == expected_issues