Skip to content
Snippets Groups Projects
test_compute_matches.py 5.01 KiB
import pytest

from nerval import evaluate

THRESHOLD = 0.30


# fmt: off
fake_tags_aligned_nested_perfect = [
    "B-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "B-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "O",
    "B-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC",
    "O",
    "O",
]


fake_tags_aligned_nested_false = [
    "B-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "O",
    "B-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC",
    "O",
    "O",
]

fake_predict_tags_aligned = [
    "B-PER", "I-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER",
    "I-PER",
    "I-PER", "I-PER", "I-PER", "I-PER", "I-PER", "I-PER", "I-PER", "I-PER",
    "O",
    "O", "O", "O", "O",
    "O", "O", "O", "O", "O",
    "O",
    "O", "O",
    "O",
    "B-***", "I-***", "I-***", "I-***", "I-***",
    "O",
    "O", "O",
    "O",
    "B-DAT", "I-DAT", "I-DAT", "I-DAT",
    "O",
    "O", "O",
]
# fmt: on

fake_annot_tags_aligned = [
    "B-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "O",
    "B-LOC",
    "I-LOC",
    "I-LOC",
    "I-LOC",
    "I-LOC",
    "O",
    "O",
    "O",
    "O",
    "B-DAT",
    "I-DAT",
    "I-DAT",
    "I-DAT",
    "O",
    "O",
    "O",
]

fake_annot_tags_bk_boundary = [
    "O",
    "O",
    "O",
    "O",
    "B-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
]

fake_predict_tags_bk_boundary = [
    "B-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
]

fake_annot_tags_bk_boundary_2 = [
    "O",
    "O",
    "B-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
]

fake_predict_tags_bk_boundary_2 = [
    "B-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
    "I-PER",
]


@pytest.mark.parametrize(
    "test_input, expected",
    [
        (
            (
                "Gérard de -N-erval was bo-rn in Paris in 1808 -.",
                "G*rard de *N*erval ----bo*rn in Paris in 1833 *.",
                fake_annot_tags_aligned,
                fake_predict_tags_aligned,
                THRESHOLD,
            ),
            {"All": 1, "PER": 1, "LOC": 0, "DAT": 0},
        ),
        (
            (
                "Louis par la grâce de Dieu roy de France et de Navarre.",
                "Louis par la grâce de Dieu roy de France et de Navarre.",
                fake_tags_aligned_nested_perfect,
                fake_tags_aligned_nested_perfect,
                THRESHOLD,
            ),
            {"All": 3, "PER": 1, "LOC": 2},
        ),
        (
            (
                "Louis par la grâce de Dieu roy de France et de Navarre.",
                "Louis par la grâce de Dieu roy de France et de Navarre.",
                fake_tags_aligned_nested_perfect,
                fake_tags_aligned_nested_false,
                THRESHOLD,
            ),
            {"All": 2, "PER": 1, "LOC": 1},
        ),
        (
            (
                "The red dragon",
                "The red dragon",
                fake_annot_tags_bk_boundary,
                fake_predict_tags_bk_boundary,
                THRESHOLD,
            ),
            {"All": 0, "PER": 0},
        ),
        (
            (
                "A red dragon",
                "A red dragon",
                fake_annot_tags_bk_boundary_2,
                fake_predict_tags_bk_boundary_2,
                THRESHOLD,
            ),
            {"All": 1, "PER": 1},
        ),
    ],
)
def test_compute_matches(test_input, expected):
    assert evaluate.compute_matches(*test_input) == expected


def test_compute_matches_empty_entry():
    with pytest.raises(AssertionError):
        evaluate.compute_matches(None, None, None, None, None)