doczyai-pipelines/fieldExtraction/tests/test_chunking.py

import pytest

from src.preprocessing_funcs import smart_chunk_ac

test_chunk_dict_term = {
    "1": "NA",
    "2": "Termination",
    "3": "NA",
    "4": "NA",
}

test_full_text_term = '\n'.join([v for v in test_chunk_dict_term.values()])

test_late_paid_claims = {
    "1": "NA",
    "2": "LATE PAYMENT",  # try lowercase
    "3": "NA",
    "4": "NA",
    "10": "NA",
    "11": "LATE PAID",
    "12": "NA",
}

test_full_text_late_paid_claims = '\n'.join([v for v in test_late_paid_claims.values()])

test_keyword_mappings = {
    "or_1": {
        "methodology": "or",
        "keywords": ["kw1", "kw2", "kw3"],
        "case_sensitive": False,
    },
    "or_2": {"methodology": "or", "keywords": ["kw4"], "case_sensitive": False},
    "hierarchy_1": {
        "methodology": "hierarchy",
        "keywords": ["h5", "h6"],
        "case_sensitive": False,
    },
    "cs_or_1": {
        "methodology": "or",
        "keywords": ["KW1", "kw2", "kw3"],
        "case_sensitive": True,
    },
    "and": {"methodology": "and", "keywords": ["KW1", "KW2"], "case_sensitive": False},
    "regex": {
        "methodology": "regex",
        "regex": r"\b\d{2}-?\d{7}\b",
        "case_sensitive": True,
        "keywords": ["KW1", "KW2"],
    },
}

case_1 = {
    "1": "NA",
    "2": "kw1",
    "3": "NA",
    "4": "h6",
    "5": "NA",
    "6": "kw3",
}

case_2 = {
    "1": "NA",
    "2": "h6",
    "3": "NA",
    "4": "kw2",
    "10": "NA",
    "11": "h5",
    "12": "NA",
    "15": "kw4",
}

case_3 = {
    "1": "NA",
    "2": "KW1",
    "3": "NA",
    "4": "h6",
    "5": "NA",
    "6": "kw3",
}

case_4 = {
    "1": "NA",
    "2": "KW1",
    "3": "NA",
    "4": "KW1KW2",
    "5": "NA",
    "6": "kw3",
}

case_5 = {
    "1": "NA",
    "2": "12-3456789",
    "3": "987654321",
    "4": "Next",
    "5": "1234567-89",
    "6": "kw3",
}


def test_smart_chunk_term():
    assert smart_chunk_ac(text_dict=test_chunk_dict_term, 
                          contract_text=test_full_text_term)["term_group"] == "\n".join(["NA", "Termination", "NA"])


def test_smart_chunk_late_paid():
    assert smart_chunk_ac(text_dict=test_late_paid_claims, 
                          contract_text=test_full_text_late_paid_claims)[
        "LATE_PAID_CLAIMS_LANGUAGE"
    ] == "\n".join(["NA", "LATE PAID", "NA"])


def test_or_1():
    assert smart_chunk_ac(text_dict=case_1, 
                          contract_text='\n'.join([v for v in case_1.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "or_1"
    ] == "\n".join(["NA", "kw1", "NA", "NA", "kw3"])
    assert smart_chunk_ac(text_dict=case_2, 
                          contract_text='\n'.join([v for v in case_2.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "or_1"
    ] == "\n".join(["NA", "kw2"])


def test_or_2():
    assert smart_chunk_ac(text_dict=case_1, 
                          contract_text='\n'.join([v for v in case_1.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "or_2"
    ] == "\n".join([])
    assert smart_chunk_ac(text_dict=case_2, 
                          contract_text='\n'.join([v for v in case_2.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "or_2"
    ] == "\n".join(["kw4"])


def test_case_sensitive_or_1():
    assert smart_chunk_ac(text_dict=case_1, 
                          contract_text='\n'.join([v for v in case_1.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "cs_or_1"
    ] == "\n".join(["NA", "kw3"])
    assert smart_chunk_ac(text_dict=case_3, 
                          contract_text='\n'.join([v for v in case_3.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "cs_or_1"
    ] == "\n".join(["NA", "KW1", "NA", "NA", "kw3"])


def test_hierarchy_1():
    assert smart_chunk_ac(text_dict=case_1, 
                          contract_text='\n'.join([v for v in case_1.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "hierarchy_1"
    ] == "\n".join(["NA", "h6", "NA"])
    assert smart_chunk_ac(text_dict=case_2, 
                          contract_text='\n'.join([v for v in case_2.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "hierarchy_1"
    ] == "\n".join(["NA", "h5", "NA"])


def test_and():
    assert smart_chunk_ac(text_dict=case_4, 
                          contract_text='\n'.join([v for v in case_4.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "and"
    ] == "\n".join(["NA", "KW1KW2", "NA"])


def test_regex():
    assert smart_chunk_ac(text_dict=case_5, 
                          contract_text='\n'.join([v for v in case_5.values()]),
                          keyword_mappings=test_keyword_mappings)[
        "regex"
    ] == "\n".join(["NA", "12-3456789", "987654321", "Next"])


def test_undefined_methodology():
    with pytest.raises(ValueError):
        smart_chunk_ac(
            text_dict=case_1,
            contract_text='\n'.join([v for v in case_1.values()]), 
            keyword_mappings={
                "e1": {
                    "methodology": "Not a real methodology",
                    "keywords": ["NA"],
                    "case_sensitive": True,
                }
            },
        )