344 lines
12 KiB
Python
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
|