Katon Minhas
8c3c674e75
Feature/daip2 33 optimize exhibit start prompt * Updated Exhibit Check prompt * Updated prompt, reorg * Switch to claude 3 haiku for new exhibit check prompt Approved-by: Alex Galarce
315 lines
13 KiB
Python
315 lines
13 KiB
Python
import re
|
|
from collections import defaultdict
|
|
from itertools import chain, count, groupby
|
|
|
|
import claude_funcs
|
|
import config
|
|
import keywords
|
|
import prompts
|
|
import smart_chunking_funcs
|
|
import string_funcs
|
|
import valid
|
|
from valid import DERIVED_INDICATOR_FIELDS
|
|
from regex_patterns import PIPE_PATTERN
|
|
|
|
|
|
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 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
|
|
|
|
|
|
# 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_funcs.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_funcs.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 = smart_chunking_funcs.chunk_or(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"]
|
|
)
|
|
keyword_page_dict = smart_chunking_funcs.keyword_search(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"]
|
|
)
|
|
elif key_dict["methodology"] == "hierarchy":
|
|
page_list = smart_chunking_funcs.chunk_hierarchical(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"]
|
|
)
|
|
keyword_page_dict = smart_chunking_funcs.keyword_search(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"]
|
|
)
|
|
elif key_dict["methodology"] == "and":
|
|
page_list = smart_chunking_funcs.chunk_and(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"]
|
|
)
|
|
keyword_page_dict = smart_chunking_funcs.keyword_search(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"]
|
|
)
|
|
elif key_dict["methodology"] == "regex":
|
|
page_list = smart_chunking_funcs.chunk_regex(text_dict, key_dict["regex"])
|
|
keyword_page_dict = smart_chunking_funcs.keyword_search(
|
|
text_dict, key_dict["keywords"], key_dict["case_sensitive"], True
|
|
)
|
|
elif key_dict["methodology"] == "include_exclude":
|
|
# Chunk pages based on included and excluded keywords
|
|
page_list = smart_chunking_funcs.chunk_with_include_exclude_keywords(
|
|
text_dict, key_dict["included_keywords"], key_dict["excluded_keywords"]
|
|
)
|
|
|
|
# Create a copy of keywords list by combining included and excluded keywords - this is required for keyword search
|
|
included_keywords = key_dict["included_keywords"]
|
|
excluded_keywords = key_dict["excluded_keywords"]
|
|
|
|
keywords = list(chain(included_keywords,excluded_keywords))
|
|
|
|
# Perform keyword search on the text dictionary - this returns a dictionary keyed by keyword and values are list of pages containing that keyword
|
|
# This is not currently used but is available for debugging purposes (or future use cases)
|
|
keyword_page_dict = smart_chunking_funcs.keyword_search(
|
|
text_dict, keywords, key_dict["case_sensitive"]
|
|
)
|
|
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):
|
|
"""
|
|
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, filename):
|
|
exhibit_pages = []
|
|
for page_num, page in text_dict.items():
|
|
prompt = prompts.EXHIBIT_CHECK(page[0:100])
|
|
claude_answer_raw = claude_funcs.invoke_claude(
|
|
prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=10
|
|
)
|
|
claude_answer_extracted = re.findall(pattern=PIPE_PATTERN, string=claude_answer_raw)[-1]
|
|
if "Y" in claude_answer_extracted:
|
|
exhibit_pages.append(page_num)
|
|
return exhibit_pages
|
|
|