Alex Galarce
f84023d742
Bugfix/table preprocessing * restructure get_exhibit_dict * exhibit header * update exhibit header handling * update docstring * bugfix * remove test - update docstring * Merged main into bugfix/table-preprocessing * refactor: update page_key_sort to return float for numeric keys * refactor: enhance get_exhibit_dict to improve chunking logic for exhibits * refactor: update page_key_sort to return float for numeric keys * test: add unit tests for get_exhibit_dict function * test: fix test case 7 * Add final, difficult test * refactor: enhance get_exhibit_dict to handle single decimal page exhibits and improve chunking logic * pulled `one_to_n_test.py` from dynamic primary branch * fix: update common_file_names logic to remove specific entry and log changes * docs: update get_exhibit_dict docstring to clarify input and output formats * test: add test for exhibit handling with decimal pages in get_exhibit_dict * feat: add debug logging for exhibit chunk mapping and exhibit dictionary keys in one_to_n_exhibit_chunking * Merge remote-tracking branch 'origin/main' into bugfix/table-preprocessing * feat: enhance spatial processing rules and consistency checks for reimbursement terms * feat: update global scan methodology to enhance context extraction for reimbursement constraints * Merge remote-tracking branch 'origin/main' into bugfix/table-preprocessing Approved-by: Katon Minhas
458 lines
19 KiB
Python
458 lines
19 KiB
Python
import re
|
|
import logging
|
|
|
|
from src.prompts import investment_prompts
|
|
import src.utils.llm_utils as llm_utils
|
|
import src.utils.string_utils as string_utils
|
|
from src import config, keywords
|
|
from src.regex.regex_patterns import PIPE_PATTERN
|
|
|
|
|
|
from src import config, keywords
|
|
from src.enums.delimiters import Delimiter
|
|
|
|
def remove_page_indicators(contract_text: str) -> str:
|
|
"""Clean textract output by removing page number indicators in the form of "Page X of Y"
|
|
|
|
This function processes input text to 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
|
|
"""
|
|
|
|
if contract_text:
|
|
cleaned_text = re.sub(r"Page [0-9]+ of [0-9]+\n\n", " ", contract_text)
|
|
else:
|
|
cleaned_text = contract_text
|
|
return cleaned_text
|
|
|
|
|
|
# TODO: write unit tests
|
|
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)):
|
|
if temp_list[i]:
|
|
text_dict[temp_list[i].split()[0]] = text_list[i] # splits on whitespace characters, which includes spaces, tabs, and newline characters
|
|
|
|
return {k: v for k, v in text_dict.items() if k != "Document"}
|
|
|
|
|
|
def clean_law_symbols(contract_text):
|
|
if not contract_text:
|
|
return contract_text
|
|
# First correction: Replace '$$' with '$'
|
|
contract_text = contract_text.replace("$$", "$")
|
|
|
|
contract_text = re.sub(r"(U\.?S\.?C\.?) \$", r"\1§", contract_text) # replaces $ with § when it follows abbreviations like "U.S.C." (United States Code).
|
|
contract_text = re.sub(r"(C\.?F\.?R\.?) \$", r"\1§", contract_text) # replaces $ with § when it follows abbreviations like "C.F.R." (Code of Federal Regulations).
|
|
|
|
# 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
|
|
|
|
|
|
# ORIGINAL
|
|
def chunk_consecutive_og(text_dict, exhibit_pages):
|
|
# If needed - extend page 1 to page 1 and 2, then cut chunking off after (edge case: compensation terms not found on first page of exhibit)
|
|
|
|
reimbursement_pages = [page_num for page_num in text_dict.keys() if string_utils.contains_reimbursement(text_dict, page_num)]
|
|
|
|
page_dict = {}
|
|
current_exhibit = None
|
|
for page_num in text_dict.keys():
|
|
|
|
# Page is Reimbursement AND Exhibit
|
|
if page_num in reimbursement_pages and page_num in exhibit_pages:
|
|
current_exhibit = page_num
|
|
page_dict[page_num] = [page_num]
|
|
|
|
# Page is Reimbursement NOT Exhibit
|
|
elif page_num in reimbursement_pages and page_num not in exhibit_pages:
|
|
if current_exhibit:
|
|
page_dict[current_exhibit].append(page_num)
|
|
else:
|
|
page_dict[page_num] = [page_num]
|
|
|
|
# Page is Exhibit NOT Reimbursement
|
|
elif page_num in exhibit_pages and page_num not in reimbursement_pages:
|
|
current_exhibit = page_num
|
|
page_dict[page_num] = [page_num]
|
|
|
|
# Page is NOT Exhibit NOT Reimbursment
|
|
elif page_num not in exhibit_pages and page_num not in reimbursement_pages:
|
|
current_exhibit = None
|
|
page_dict[page_num] = [page_num]
|
|
|
|
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]
|
|
return final_dict
|
|
|
|
|
|
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 string_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 filter_quick_review(text_dict):
|
|
"""
|
|
Cover Sheets (aka Top Sheets or Quick Review pages) are pages stapled to the front of the contract that contain manually written summaries of the contract's contents.
|
|
They are NOT legally binding documents, and because they are often manually filled out and hand-written, are more prone to have erroneous information than the rest of the contract.
|
|
For Doczy.AI, we should not pull any information from top sheets, except as a very last resort.
|
|
|
|
filter_quick_review() splits the input dictionary into two dictionaries, one containing the contract pages, the other containing the top sheet pages.
|
|
|
|
The function checks each page's text for the keywords "QUICK REVIEW", "TOP SHEET", and "COVER SHEET". It categorizes the text into two separate dictionaries: one for texts that do not contain any of these keywords, and another for texts that do.
|
|
|
|
Parameters:
|
|
text_dict (dict): A dictionary where the key is the page number and the value is the text of that page.
|
|
|
|
Returns:
|
|
tuple of two dicts:
|
|
- The first dictionary contains pages that do not have the specified keywords.
|
|
- The second dictionary includes pages that contain any of the specified keywords.
|
|
"""
|
|
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()
|
|
}, {
|
|
page_num: page_text
|
|
for page_num, page_text in text_dict.items()
|
|
if "QUICK REVIEW" in page_text.upper()
|
|
or "TOP SHEET" in page_text.upper()
|
|
or "COVER SHEET" in page_text.upper()
|
|
}
|
|
|
|
def get_exhibit_pages(text_dict: dict[str, str], filename: str) -> list[str]:
|
|
"""Extract beginning-of-exhibit pages from a contract using LLM-based detection.
|
|
|
|
This function identifies pages that start new exhibits by checking each page (or
|
|
subpage) for exhibit headers using an LLM prompt. It handles both regular pages
|
|
(e.g., "5") and subpages created by table splitting (e.g., "5.1", "5.2").
|
|
|
|
Key behaviors:
|
|
- The first page of the contract is always included in the exhibit pages.
|
|
- For base pages with subpages, if ANY subpage contains an exhibit header,
|
|
the FIRST subpage is added to the exhibit list
|
|
- Each base page number is only processed once to avoid duplicates.
|
|
|
|
Args:
|
|
text_dict (dict[str, str]): Dictionary keyed by string-formatted page numbers and valued by page text
|
|
Page numbers may include subpages (e.g., "5.1", "5.2")
|
|
filename (str): Filename of the contract (for tracking purposes)
|
|
|
|
Returns:
|
|
list[str]: Ordered list of page numbers that start new exhibits. May include
|
|
subpage numbers (e.g., ["1", "5.1", "5.2"]) depending on table splitting.
|
|
|
|
Example:
|
|
Given pages ["1", "2", "5.0", "5.1", "5.2", "7"] where "5.1" contains an exhibit header:
|
|
Returns ["1", "5.1", "7"] (assuming pages 1 and 7 also have exhibit headers)
|
|
"""
|
|
exhibit_pages = []
|
|
first_page = True
|
|
|
|
# Sort pages to handle subpages properly (5.0, 5.1, 5.2, etc.)
|
|
# TODO: augment `string_utils.page_key_sort` to handle subpages
|
|
sorted_pages = sorted(text_dict.keys(), key=lambda x: (
|
|
int(x.split('.')[0]),
|
|
float(x.split('.')[1]) if '.' in x else 0)
|
|
)
|
|
|
|
for page_num in sorted_pages:
|
|
|
|
# Always include the first page as an exhibit
|
|
if first_page:
|
|
exhibit_pages.append(page_num)
|
|
is_exhibit = True
|
|
first_page = False
|
|
else:
|
|
if "." not in page_num or ".0" in page_num:
|
|
page_content = text_dict[page_num]
|
|
prompt = investment_prompts.EXHIBIT_CHECK(page_content[0:200])
|
|
claude_answer_raw = llm_utils.invoke_claude(
|
|
prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=150
|
|
)
|
|
claude_answer_extracted = string_utils.extract_text_from_delimiters(
|
|
claude_answer_raw, Delimiter.PIPE
|
|
)
|
|
if "Y" in claude_answer_extracted:
|
|
exhibit_pages.append(page_num)
|
|
is_exhibit = True
|
|
else:
|
|
is_exhibit = False
|
|
elif is_exhibit==True:
|
|
exhibit_pages.append(page_num)
|
|
return exhibit_pages
|
|
|
|
def chunk_by_exhibit(text_dict: dict,
|
|
exhibit_pages: list
|
|
) -> dict:
|
|
"""
|
|
Organizes pages into groups based on their association with specific exhibits.
|
|
|
|
This function assigns each page number from the `text_dict` dictionary to an exhibit
|
|
based on the `exhibit_pages` list. Pages are grouped under the nearest preceding
|
|
page number in `exhibit_pages`. If a page number in `text_dict` is itself in
|
|
`exhibit_pages`, it starts a new exhibit group.
|
|
|
|
Parameters:
|
|
text_dict (dict): A dictionary where keys are page numbers and values are page text
|
|
exhibit_pages (list): A list of page numbers that mark the beginning of a new
|
|
exhibit.
|
|
|
|
Returns:
|
|
dict: A dictionary mapping each page number in `text_dict` to its corresponding
|
|
exhibit identifier. The exhibit identifier is the page number of the first
|
|
page in that exhibit as listed in `exhibit_pages`. If there are pages before
|
|
the first `exhibit_page`, they are grouped under the exhibit identifier "0".
|
|
"""
|
|
if len(exhibit_pages) == 0:
|
|
return {key : key for key in text_dict.keys()}
|
|
|
|
exhibit_chunk_mapping = {}
|
|
current_exhibit = "0"
|
|
for page_num in text_dict.keys():
|
|
if page_num in exhibit_pages:
|
|
current_exhibit = page_num
|
|
exhibit_chunk_mapping[page_num] = current_exhibit
|
|
else:
|
|
exhibit_chunk_mapping[page_num] = current_exhibit
|
|
return exhibit_chunk_mapping
|
|
|
|
|
|
def get_exhibit_dict(text_dict: dict,
|
|
exhibit_chunk_mapping: dict
|
|
) -> dict:
|
|
"""
|
|
Creates a dictionary of exhibit texts, concatenating pages with the same exhibit number,
|
|
but separating decimal pages into their own chunks.
|
|
|
|
Args:
|
|
text_dict (dict): Dictionary mapping page numbers to page text.
|
|
exhibit_chunk_mapping (dict): Dictionary mapping page numbers to exhibit numbers.
|
|
|
|
Returns:
|
|
dict: Dictionary mapping exhibit chunk keys to combined text. Format for keys:
|
|
"{exhibit}.{index}"
|
|
NOTE: If exhibit identifiers contain decimals (e.g., "23.0") output keys will
|
|
follow the same pattern (resulting in "23.0.0", "23.0.1", etc.)
|
|
|
|
Guidelines:
|
|
- Every exhibit chunk contains at most 1 table page in it (i.e. one page with a decimal)
|
|
- Tables are inserted in chronological order within the document flow.
|
|
- Multiple decimal pages are separated into individual chunks
|
|
- Each chunk gets a new key: "{exhibit}.{index}" (e.g. "2.0", "2.1", "2.2", etc.)
|
|
|
|
Algorithm:
|
|
For each exhibit:
|
|
1. Identify base pages (no decimals) and table pages (with decimals).
|
|
2. For each decimal page, create a chunk containing:
|
|
- All base pages that came before the decimal page (in document order)
|
|
- The specific decimal page
|
|
- All base pages that came after the decimal page (in document order)
|
|
|
|
Example 1:
|
|
Input mapping {'1.0': '1.0', '2': '2', '3.0': '2', '3.1': '2', '3.2': '2', '4.0': '2'}
|
|
- '1.0': text from '1.0' (single page exhibit)
|
|
- '2.0': text from '2' + text from '3.0' (base page 2, then table 3.0)
|
|
- '2.1': text from '2' + text from '3.1' (base page 2, then table 3.1)
|
|
- '2.2': text from '2' + text from '3.2' (base page 2, then table 3.2)
|
|
- '2.3': text from '2' + text from '4.0' (base page 2, then table 4.0)
|
|
|
|
Example 2:
|
|
Input mapping {'1.0': '1.0', '2': '2', '3.0': '2', '3.1': '2', '3.2': '2', '4': '2', '5' : '2'}
|
|
- '1.0': text from '1.0'
|
|
- '2.0': text from '2' + text from '3.0' + text from '4' + text from '5'
|
|
- '2.1': text from '2' + text from '3.1' + text from '4' + text from '5'
|
|
- '2.2': text from '2' + text from '3.2' + text from '4' + text from '5'
|
|
- '2.3': text from '2' + text from '4' + text from '5'
|
|
|
|
Note: Tables 3.0, 3.1, and 3.2 appear between pages 2 and 4 in document flow.
|
|
|
|
Other cases:
|
|
- No decimal pages: creates single chunk for the exhibit.
|
|
- Single decimal page exhibit: uses original page as the chunk key.
|
|
"""
|
|
exhibit_dict = {}
|
|
|
|
# Group pages by exhibit
|
|
exhibit_pages = {}
|
|
for page, exhibit in exhibit_chunk_mapping.items():
|
|
if exhibit not in exhibit_pages:
|
|
exhibit_pages[exhibit] = []
|
|
exhibit_pages[exhibit].append(page)
|
|
|
|
# Process each exhibit
|
|
for exhibit in exhibit_pages:
|
|
pages = sorted(exhibit_pages[exhibit], key=string_utils.page_key_sort)
|
|
|
|
base_pages = [p for p in pages if '.' not in p]
|
|
decimal_pages = [p for p in pages if '.' in p]
|
|
|
|
# Special case: single decimal page exhibit
|
|
if len(decimal_pages) == 1 and not base_pages:
|
|
exhibit_dict[decimal_pages[0]] = text_dict[decimal_pages[0]]
|
|
continue
|
|
|
|
# Case: No decimal pages - create single chunk with .0 suffix
|
|
if not decimal_pages:
|
|
exhibit_dict[f"{exhibit}.0"] = "\n".join(text_dict[p] for p in pages)
|
|
continue
|
|
|
|
# Case: With decimal pages, create chunks for each decimal page
|
|
chunk_index = 0
|
|
for decimal_page in decimal_pages:
|
|
chunk_key = f"{exhibit}.{chunk_index}"
|
|
|
|
# Get decimal page base number for chronological positioning
|
|
decimal_base = int(decimal_page.split('.')[0])
|
|
|
|
# Build chunk with pages in chronological order
|
|
chunk_parts = []
|
|
|
|
# Add base pages before decimal page
|
|
for base_page in base_pages:
|
|
if int(base_page) <= decimal_base:
|
|
chunk_parts.append(text_dict[base_page])
|
|
|
|
# Add the decimal page itself
|
|
chunk_parts.append(text_dict[decimal_page])
|
|
|
|
# Add base pages after decimal page
|
|
for base_page in base_pages:
|
|
if int(base_page) > decimal_base:
|
|
chunk_parts.append(text_dict[base_page])
|
|
|
|
exhibit_dict[chunk_key] = "\n".join(chunk_parts)
|
|
chunk_index += 1
|
|
|
|
# Add final chunk with just base pages (if there are base pages after all decimal pages)
|
|
if base_pages and decimal_pages:
|
|
# Check if there are base pages that come after all decimal pages
|
|
max_decimal_base = max(int(dp.split('.')[0]) for dp in decimal_pages)
|
|
remaining_base_pages = [bp for bp in base_pages if int(bp) > max_decimal_base]
|
|
|
|
if remaining_base_pages:
|
|
chunk_key = f"{exhibit}.{chunk_index}"
|
|
# Include ALL base pages in the final chunk
|
|
exhibit_dict[chunk_key] = "\n".join(text_dict[bp] for bp in base_pages)
|
|
return exhibit_dict
|
|
|
|
def get_exhibit_headers(exhibit_dict, filename):
|
|
"""
|
|
Extracts the exhibit header using an LLM prompt.
|
|
|
|
Args:
|
|
exhibit_chunk (str): A chunk of exhibit text to analyze.
|
|
filename (str): The name of the file being processed.
|
|
|
|
Returns:
|
|
str: The extracted exhibit header.
|
|
"""
|
|
exhibit_header_dict = {}
|
|
base_pages = {}
|
|
for page_num, exhibit_chunk in exhibit_dict.items():
|
|
base_page_num = page_num.split('.')[0]
|
|
if base_page_num in base_pages.keys():
|
|
exhibit_header_dict[page_num] = base_pages[base_page_num]
|
|
else:
|
|
prompt = investment_prompts.EXHIBIT_HEADER(exhibit_chunk[0:400])
|
|
llm_answer_raw = llm_utils.invoke_claude(
|
|
prompt, config.MODEL_ID_CLAUDE35_SONNET, filename
|
|
)
|
|
llm_answer_final = string_utils.extract_text_from_delimiters(
|
|
llm_answer_raw, Delimiter.PIPE
|
|
)
|
|
exhibit_header_dict[page_num] = llm_answer_final
|
|
base_pages[base_page_num] = llm_answer_final
|
|
|
|
return exhibit_header_dict |