Merged in feature/exhibit-linking (pull request #564)
Feature/exhibit linking * Initialize new exhibit header prompt * Save * Merge branch 'main' into feature/exhibit-linking * update prompt * cleanup return type * remove test * Test fixes * Fix test * Add tests * Merge branch 'main' into feature/exhibit-linking * Merge branch 'main' into feature/exhibit-linking * Remove test notebooks * Delete test notebook * Save * Merge branch 'main' into feature/exhibit-linking * try except for exhbit page finding * Prep for merge * update test * Remove test * Merge branch 'main' into feature/exhibit-linking * update unit test * Fix unit test * replace try except * Update page sort * Add docstring Approved-by: Alex Galarce
This commit is contained in:
@@ -211,67 +211,108 @@ def filter_quick_review(text_dict):
|
||||
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)
|
||||
def get_exhibit_pages(text_dict: dict[str, str], filename: str) -> tuple[list[str], dict[str, str]]:
|
||||
"""
|
||||
Identify exhibit pages from the text dictionary.
|
||||
This function processes the text dictionary to identify pages that are considered exhibits.
|
||||
It uses the first page as an exhibit and checks subsequent pages for exhibit headers.
|
||||
If a page is identified as an exhibit, it is added to the exhibit pages list and its header is stored in the exhibit header dictionary.
|
||||
If a page does not have a header or is not an exhibit, it is skipped.
|
||||
Args:
|
||||
text_dict (dict[str, str]): A dictionary where keys are page numbers (as strings) and values are the text of those pages.
|
||||
filename (str): The name of the file being processed, used for logging and LLM invocation.
|
||||
Returns:
|
||||
tuple[list[str], dict[str, str]]: A tuple containing:
|
||||
- A list of page numbers (as strings) that are identified as exhibit pages.
|
||||
- A dictionary where keys are page numbers and values are the headers of those exhibits.
|
||||
"""
|
||||
sorted_pages = sorted(text_dict.keys(), key=string_utils.page_key_sort)
|
||||
exhibit_pages = []
|
||||
exhibit_header_dict = {}
|
||||
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:
|
||||
|
||||
page_content = text_dict[page_num]
|
||||
# Always include the first page as an exhibit
|
||||
if first_page:
|
||||
exhibit_pages.append(page_num)
|
||||
exhibit_header_dict[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])
|
||||
prompt = investment_prompts.EXHIBIT_HEADER(page_content[0:400])
|
||||
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:
|
||||
if "N/A" in claude_answer_extracted:
|
||||
is_exhibit = False
|
||||
else:
|
||||
exhibit_pages.append(page_num)
|
||||
exhibit_header_dict[page_num] = claude_answer_extracted
|
||||
is_exhibit = True
|
||||
elif is_exhibit==True:
|
||||
exhibit_pages.append(page_num)
|
||||
return exhibit_pages
|
||||
return exhibit_pages, exhibit_header_dict
|
||||
|
||||
|
||||
def link_exhibit_pages(all_exhibit_headers, filename):
|
||||
"""
|
||||
Filters exhibit pages to include only those with meaningfully different headers.
|
||||
This function processes exhibit headers sequentially and uses literal comparison and
|
||||
LLM-based semantic comparison to determine when exhibits represent continuation pages
|
||||
vs. distinct exhibits. Continuation pages are filtered out of the headers list and
|
||||
mapping, and when used downstream are included in their parent exhibit's text chunk.
|
||||
Links exhibit pages based on their headers. If the header of a page is identical (either literally or semantically) to the previous page's header, it is not added to the final list of exhibit pages.
|
||||
This function will link contiguous header pages together, so that the final list of exhibit pages will only contain pages that are meaningfully different from the previous page. E.g. "1" with header {"Exhibit 1"} and "2" with header {"Continued exhibit 1"} will be linked together, but "3" with header {"Exhibit 2"} will not be linked to either of them.
|
||||
Args:
|
||||
all_exhibit_headers (dict[str, str]): A dictionary mapping page numbers (as strings) to their exhibit headers. Expected to be in chronological order
|
||||
filename (str): The name of the file being processed, used for logging and LLM invocation.
|
||||
Returns:
|
||||
tuple[list[str], dict[str, str]]: A tuple containing:
|
||||
- A list of page numbers (as strings) that are identified as exhibit pages.
|
||||
- A dictionary where keys are page numbers and values are the headers of those exhibits.
|
||||
Example:
|
||||
Input: {"1": "Exhibit A", "2": "Exhibit A (continued)", "3": "Exhibit B"}
|
||||
Output (["1", "3"], {"1": "Exhibit A", "3": "Exhibit B"})
|
||||
Notes:
|
||||
- First exhibit page is always included.
|
||||
- Literal header comparison is performed first to optimize LLM usage and latency.
|
||||
- Uses `EXHIBIT_LINKAGE` prompt for semantic similarity assessment
|
||||
- Expects LLM to return "Y" or "N" enclosed in pipe delimiters
|
||||
- Pages with non-distinct headers are excluded from the final results.
|
||||
"""
|
||||
first_exhibit = True
|
||||
final_exhibit_pages = []
|
||||
final_exhibit_headers = {}
|
||||
for page_num, page_header in all_exhibit_headers.items():
|
||||
if first_exhibit:
|
||||
exhibit_is_different = "N"
|
||||
previous_header = page_header
|
||||
final_exhibit_pages.append(page_num)
|
||||
final_exhibit_headers[page_num] = page_header
|
||||
first_exhibit = False
|
||||
else:
|
||||
# If literally identical, skip prompt
|
||||
if page_header.upper().strip() == previous_header.upper().strip():
|
||||
exhibit_is_different = "N"
|
||||
else:
|
||||
prompt = investment_prompts.EXHIBIT_LINKAGE(page_header, previous_header)
|
||||
claude_answer_raw = llm_utils.invoke_claude(prompt, config.MODEL_ID_CLAUDE35_SONNET, filename)
|
||||
exhibit_is_different = string_utils.extract_text_from_delimiters(claude_answer_raw, Delimiter.PIPE)
|
||||
|
||||
# If the prompt has identified that the current exhibit is meaningfully different than the previous
|
||||
if "Y" in exhibit_is_different:
|
||||
final_exhibit_pages.append(page_num)
|
||||
final_exhibit_headers[page_num] = page_header
|
||||
|
||||
previous_header = page_header
|
||||
|
||||
return final_exhibit_pages, final_exhibit_headers
|
||||
|
||||
|
||||
def chunk_by_exhibit(text_dict: dict,
|
||||
exhibit_pages: list
|
||||
|
||||
Reference in New Issue
Block a user