doczyai-pipelines/fieldExtraction/src/preprocessing_funcs.py

from itertools import groupby, count
import re
from collections import defaultdict
import utils
import valid
from valid import DERIVED_INDICATOR_FIELDS
import keywords
from collections import defaultdict
import ac_smart_chunking
import ac_funcs
import prompts
import claude_funcs
import config


def clean_newlines(contract_text: str) -> str:
    """Clean textract output by removing newlines and page number indicators.

    This function processes input text to:
    1. Remove all newline characters.
    2. Remove lines that indicate page numbers (e.g. 'Page 1 of 10')

    Args:
        contract_text (str): Raw text output from Textract to be cleaned

    Returns:
        str: cleaned text with newlines and page number indicators removed
    """
    cleaned_text = re.sub(r"(?<!\n)\n(?!\n)", " ", contract_text)
    cleaned_text = re.sub(r"Page [0-9]+ of [0-9]+\n\n", " ", contract_text)
    return cleaned_text


def split_text(text: str) -> dict[str, str]:
    """Split text on pages by the string `Start of Page No. = '

    Args:
        text (str): Raw text output from Textract to be split

    Returns:
        dict[str, str]: A dictionary, keyed by the string page number and valued by the page text.
    """
    temp_list = text.split("Start of Page No. = ")
    text_list = re.split(r"Start of Page No. = [0-9]+\n", text)

    text_dict = {}
    for i in range(len(text_list)):
        text_dict[temp_list[i].split()[0]] = text_list[i]

    return {k: v for k, v in text_dict.items() if k != "Document"}


def highlight_rates(text_dict):
    for page, text in text_dict.items():
        words = text.split()
        highlighted_words = []
        for word in words:
            if "%" in word or "$" in word:
                word = f">>>{word}<<<"
            highlighted_words.append(word)
        text_dict[page] = " ".join(highlighted_words)
    return text_dict


def clean_law_symbols(contract_text):
    contract_text = contract_text.replace("$$", "$")

    contract_text = re.sub(r"(U\.?S\.?C\.?) \$", r"\1§", contract_text)
    contract_text = re.sub(r"(C\.?F\.?R\.?) \$", r"\1§", contract_text)

    # Second correction: Replace '$' with '§' when followed by a number with three decimal places
    contract_text = re.sub(r"\$(?=\d+\.\d{3})", "§", contract_text)

    return contract_text



def chunk_consecutive(text_dict, exhibit_pages):
    exhibit_pages = set(str(page) for page in exhibit_pages)
    reimbursement_pages = {
        page_num
        for page_num in text_dict.keys()
        if utils.contains_reimbursement(text_dict, page_num)
    }

    page_dict = {}
    current_chunk_start = None
    last_exhibit = None
    in_exhibit = False

    def word_count(text):
        return len(text.split())

    for page_str in sorted(text_dict.keys(), key=int):
        page_num = int(page_str)

        if page_str in exhibit_pages:
            current_chunk_start = page_str
            last_exhibit = page_str
            in_exhibit = True
            page_dict[current_chunk_start] = [page_str]
            # print(f"found page {page_num} in exhibit, starting new chunk")

        if page_str in reimbursement_pages:
            if not in_exhibit or current_chunk_start is None:
                # Check if this page has less than 200 words and should be added to the previous chunk
                if current_chunk_start and word_count(text_dict[page_str]) < 200:
                    page_dict[current_chunk_start].append(page_str)
                    # print(f"found page {page_num} in reimbursement with less than 200 words, adding to previous chunk")
                else:
                    current_chunk_start = page_str
                    page_dict[current_chunk_start] = [page_str]
                    # print(f"found page {page_num} in reimbursement, starting new chunk")
            else:
                page_dict[current_chunk_start].append(page_str)
                # print(f"found page {page_num} in reimbursement, adding to current chunk")
        else:
            if in_exhibit and current_chunk_start is not None:
                page_dict[current_chunk_start].append(page_str)
                # print(f"found non-reimbursement page {page_num} in exhibit, adding to current chunk")
            else:
                in_exhibit = False

        # Check next page if it's non-reimbursement and not in exhibit
        next_page_str = str(page_num + 1)
        if (
            next_page_str in text_dict
            and next_page_str not in reimbursement_pages
            and next_page_str not in exhibit_pages
        ):
            if current_chunk_start is not None:
                page_dict[current_chunk_start].append(next_page_str)
                # print(f'next page {next_page_str} was found to be a non-reimbursement and added')

        # If we've moved past the last exhibit page, reset in_exhibit
        if in_exhibit and int(page_str) > int(last_exhibit):
            in_exhibit = False

    # Remove duplicates and sort page numbers in each chunk
    for key in page_dict:
        page_dict[key] = sorted(list(set(page_dict[key])), key=int)

    final_dict = {page_num: "" for page_num in page_dict.keys()}
    for page_num in page_dict.keys():
        for p in page_dict[page_num]:
            final_dict[page_num] += text_dict[p]

    # print("Final page_dict:", {k: v for k, v in page_dict.items()})
    return final_dict


def smart_chunk_ac(
    text_dict: dict, 
    contract_text: str,
    keyword_mappings: dict = keywords.GROUPED_KEYWORD_MAPPINGS,
) -> dict:
    """Performs a keyword search on textract outputs to retrieve relevant chunks for a given field.

    Args:
        text_dict (dict): Dictionary keyed by string page-number and valued by actual textract output page
        contract_text(str): Contract Text

    Raises:
        ValueError: When the `key_dict` from keyword_mappings has a methodology that is not implemented.
                    Current options are `or` and `hierarchy`.

    Returns:
        dict: A dictionary keyed by field (defined by the keys in keyword_mappings) valued by
              corresponding chunks of text (input strings delimited by newlines).
    """
    field_pages = defaultdict(set)  # TODO: remove this, it's unused
    chunks = defaultdict(str)

    full_context = "\n".join(
        [f"Page {page}:\n{content}" for page, content in text_dict.items()]
    )

    for field_group, key_dict in keyword_mappings.items():
        if key_dict["methodology"] == "or":
            page_list = ac_smart_chunking.chunk_or(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
            keyword_page_dict = ac_smart_chunking.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
        elif key_dict["methodology"] == "hierarchy":
            page_list = ac_smart_chunking.chunk_hierarchical(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
            keyword_page_dict = ac_smart_chunking.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
        elif key_dict["methodology"] == "and":
            page_list = ac_smart_chunking.chunk_and(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
            keyword_page_dict = ac_smart_chunking.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
        elif key_dict["methodology"] == "regex":
            page_list = ac_smart_chunking.chunk_regex(text_dict, key_dict["regex"])
            keyword_page_dict = ac_smart_chunking.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"], True
            )
        else:
            page_list = list(text_dict.keys())
            keyword_page_dict = {"All pages": list(text_dict.keys())}
            raise ValueError(
                "methodology must be 'hierarchy' or 'or' - using entire contract as chunk"
            )
        if field_group == "initial_page_group":
            keyword_page_dict["initial_pages"] = [1, 2]
        
        # chunks[field_group] = concatenated page texts
        if len(page_list) == len(text_dict.keys()):
            chunks[field_group] = contract_text
        else:
            chunks[field_group] = "\n".join(
                [text_dict[str(page_num)] for page_num in page_list]
            )

        chunks[field_group + "_pages"] = keyword_page_dict
        chunks[field_group + "_methodology"] = key_dict["methodology"]
        chunks[field_group + "_case"] = key_dict["case_sensitive"]
    return chunks


def filter_quick_review(text_dict):
    return {
        page_num: page_text
        for page_num, page_text in text_dict.items()
        if "QUICK REVIEW" not in page_text.upper()
        and "TOP SHEET" not in page_text.upper()
        and "COVER SHEET" not in page_text.upper()
    }