doczyai-pipelines/src/pipelines/shared/preprocessing/preprocess.py

import logging
from typing import TYPE_CHECKING

from src import config
from src.pipelines.shared.preprocessing import preprocessing_funcs

if TYPE_CHECKING:
    from src.pipelines.shared.extraction.page_funcs import Page


def clean_text(contract_text):

    contract_text = preprocessing_funcs.remove_page_indicators(contract_text)
    contract_text = preprocessing_funcs.clean_law_symbols(contract_text)

    return contract_text


def split_text(contract_text):
    """
    Splits the contract text into separate pages and filters pages needing quick review.

    This function processes the input contract text by splitting it into individual pages
    and then categorizes these pages. Pages representing a Quick Review or Cover Page section are separated from
    the main text.

    Parameters:
        contract_text (str): The full text of the contract to be processed.

    Returns:
        tuple: A tuple containing three elements:
            - dict: A dictionary with keys as page numbers (str) and values as the text of those pages.
            - dict: A dictionary of pages that were identified Quick Review.
            - int: The total number of pages in the contract text.

    Notes:
        - This function utilizes `preprocessing_funcs.split_text` to divide the contract text
          into a dictionary with page numbers as keys.
        - It uses `preprocessing_funcs.filter_quick_review` to separate out pages that need quick
          review from the main text dictionary.
    """
    text_dict = preprocessing_funcs.split_text(
        contract_text
    )  # return a dictionary with keys - page_num (str), values as the page_text
    text_dict, top_sheet_dict = preprocessing_funcs.filter_quick_review(text_dict)
    return text_dict, top_sheet_dict


def split_text_with_pages(
    contract_text: str,
) -> tuple[dict[str, "Page"], dict[str, str]]:
    """
    Split text into pages and convert to Page objects with table splitting.

    This function extends split_text() by converting the text dictionary to Page objects
    and splitting large tables into sub-pages.

    Args:
        contract_text: The full text of the contract to be processed

    Returns:
        tuple: (pages_dict, top_sheet_dict) where:
            - pages_dict: Dictionary mapping page numbers to Page objects
            - top_sheet_dict: Dictionary of pages identified as Quick Review
    """
    # First, split into text_dict (existing logic)
    text_dict, top_sheet_dict = split_text(contract_text)

    # Then, convert to Page objects and split large tables
    pages_dict = preprocessing_funcs.split_large_tables(text_dict)

    return pages_dict, top_sheet_dict


def one_to_n_exhibit_chunking(
    text_dict=None,
    pages_dict=None,
    EXHIBIT_HEADER_MARKERS=None,
    filename=None,
) -> tuple[dict, dict]:
    """
    Process the text dictionary or pages dictionary to identify and chunk exhibits.

    This function extracts exhibit pages, maps them to their respective chunks,
    and identifies reimbursement exhibits using LLM.

    Args:
        text_dict: A dictionary where keys are page numbers and values are the text
                  on those pages (for backward compatibility)
        pages_dict: A dictionary where keys are page numbers and values are Page objects
                   (preferred)
        EXHIBIT_HEADER_MARKERS: List of markers to identify exhibit headers
        filename: The name of the file being processed.

    Returns:
        tuple: A tuple containing:
            - dict: A dictionary where keys are exhibit page numbers and values are
                   lists of page numbers that belong to that exhibit.
            - dict: A dictionary where keys are exhibit page numbers and values are
                   the exhibit headers.

    Note:
        Either text_dict or pages_dict must be provided. If both are provided,
        pages_dict takes precedence.
    """
    # Determine which dictionary to use
    if pages_dict is not None:
        exhibit_pages, all_exhibit_headers = preprocessing_funcs.get_exhibit_pages(
            pages_dict=pages_dict,
            EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS,
            filename=filename,
        )
        exhibit_pages, all_exhibit_headers = preprocessing_funcs.link_exhibit_pages(
            all_exhibit_headers, pages_dict=pages_dict, filename=filename
        )
        exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit(
            pages_dict=pages_dict, exhibit_pages=exhibit_pages
        )
    elif text_dict is not None:
        exhibit_pages, all_exhibit_headers = preprocessing_funcs.get_exhibit_pages(
            text_dict=text_dict,
            EXHIBIT_HEADER_MARKERS=EXHIBIT_HEADER_MARKERS,
            filename=filename,
        )
        exhibit_pages, all_exhibit_headers = preprocessing_funcs.link_exhibit_pages(
            all_exhibit_headers, text_dict=text_dict, filename=filename
        )
        exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit(
            text_dict=text_dict, exhibit_pages=exhibit_pages
        )
    else:
        raise ValueError("Either text_dict or pages_dict must be provided")

    return exhibit_chunk_mapping, all_exhibit_headers


def find_headers_and_footers(text_dict: dict[str, str]) -> tuple[dict, str, str]:
    """
    Returns a tuple of two lists: (header_markers, footer_markers).

    header_markers: The longest substring that is found on EVERY page of the text_dict at the START of the page.
    footer_markers: The longest substring that is found on EVERY page of the text_dict at the END of the page.

    Args:
        text_dict (dict): A dictionary where keys are page numbers and values are the text of those pages.
    Returns:
        tuple: A tuple containing two strings:
            - header_marker: A string representing the longest common substring found at the start of each page. Can be empty string
            - footer_marker: A string representing the longest common substring found at the foot of each page. Can be empty string
    """

    if not text_dict:
        return {}, "", ""

    # Get the pages as a list
    pages = list(text_dict.values())

    # If there's only one page, return empty strings and original dict
    if len(pages) <= 1:
        return text_dict.copy(), "", ""

    # Split pages into words
    pages_words = [page.split() for page in pages]

    # Find header (common prefix words)
    header_words = []
    min_word_count = min(len(page_words) for page_words in pages_words)

    for i in range(min_word_count):
        word = pages_words[0][i]
        if all(page_words[i] == word for page_words in pages_words[1:]):
            header_words.append(word)
        else:
            break

    # Find footer (common suffix words)
    footer_words = []
    for i in range(1, min_word_count + 1):
        word = pages_words[0][-i]
        if all(page_words[-i] == word for page_words in pages_words[1:]):
            footer_words.insert(0, word)  # Insert at beginning to maintain order
        else:
            break

    # Convert back to strings with proper spacing
    header_marker = " ".join(header_words)
    footer_marker = " ".join(footer_words)

    # Create a modified text_dict with headers and footers removed
    modified_text_dict = {}

    first_page = True
    for page_num, page_text in text_dict.items():
        modified_text = page_text

        if not first_page:  # Keep global header for the first page that it occurs
            # Remove header if it exists
            if header_marker:
                modified_text = modified_text[len(header_marker) :].lstrip()

            # Remove footer if it exists
            if footer_marker:
                modified_text = modified_text[: -len(footer_marker)].rstrip()

            modified_text_dict[page_num] = modified_text
        else:
            first_page = False
            modified_text_dict[page_num] = page_text

    return modified_text_dict, header_marker, footer_marker