From 3bb6abbcefcd3b6c4b274e69a4fcb59ebc23ce7e Mon Sep 17 00:00:00 2001 From: Katon Minhas Date: Fri, 10 Jan 2025 17:24:44 +0000 Subject: [PATCH] Merged in feature/simplify-b-field-chunking (pull request #348) feature/simplify b field chunking * Add chunk_by_exhibit function and unit tests * Refactor preprocess * Updated to run BU primary on entire Exhibit * Makes BU secondary compatible with exhibit chunk mapping * Added unit tests for get_exhibit_chunk * Trimmed TD funcs * Downstream changes to run_top_down * Removed conditional funcs * Merged main into feature/simplify-b-field-chunking * removed conditional funcs import * Merged main into feature/simplify-b-field-chunking * Fixed run_bottom_up_primary arguments * Fixed TD primary input * Fixed grouped keyword mappings * Remove filename tracking until postprocessing * Removed postprocessing for client columns * Removed col-based postprocessing * Merged main into feature/simplify-b-field-chunking * Fixed NoneType * Add docstrings * Updated unit tests for removal of 'Filename' * Updated comments Approved-by: Alex Galarce --- .../src/investment/conditional_funcs.py | 104 -------- .../src/investment/file_processing.py | 39 +-- .../src/investment/one_to_n_funcs.py | 241 ++++-------------- fieldExtraction/src/investment/postprocess.py | 38 +-- fieldExtraction/src/investment/preprocess.py | 99 +++++-- fieldExtraction/src/preprocessing_funcs.py | 36 +++ fieldExtraction/src/utils/string_utils.py | 29 ++- .../tests/exhibit_chunk_mapping_test.py | 31 +++ .../tests/get_exhibit_chunk_test.py | 34 +++ fieldExtraction/tests/string_utils_test.py | 18 +- 10 files changed, 278 insertions(+), 391 deletions(-) delete mode 100644 fieldExtraction/src/investment/conditional_funcs.py create mode 100644 fieldExtraction/tests/exhibit_chunk_mapping_test.py create mode 100644 fieldExtraction/tests/get_exhibit_chunk_test.py diff --git a/fieldExtraction/src/investment/conditional_funcs.py b/fieldExtraction/src/investment/conditional_funcs.py deleted file mode 100644 index 5a12655..0000000 --- a/fieldExtraction/src/investment/conditional_funcs.py +++ /dev/null @@ -1,104 +0,0 @@ -import src.utils.llm_utils as llm_utils -from src import config, prompts - - -def run_conditional(combined_results, text_dict, filename): - - dsh_dict, rbrvs_dict, sl_dict = {}, {}, {} - - for d in combined_results: - if d is not None: - page_num = d["page_num"] - - # LOB/Program Check - if "," in d["CONTRACT_LOB"]: - prompt = prompts.CONDITIONAL_LOB_CHECK(d, text_dict[page_num]) - lob_answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=128 - ) - d["CONTRACT_LOB"] = lob_answer - - # Prov Type Level 2 - prompt = prompts.CONDITIONAL_PROV_2(d, text_dict[page_num]) - prov2_answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=128 - ) - d["PROV_TYPE_LEVEL_2"] = prov2_answer.strip() - - # # Inclusion of essential RBRVS "Fee Source" Language (Y/N) - Top Down approach to be used for non-NY - if "NY" not in config.STATE: - prompt = prompts.CONDITIONAL_RBRVS(text_dict[page_num]) - rbrvs_answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE2, filename, max_tokens=16 - ) - - d['Inclusion of essential RBRVS "Fee Source" Language (Y/N)'] = ( - rbrvs_answer.strip() - ) - rbrvs_dict[page_num] = rbrvs_answer - else: - if "RBRVS" in d["FULL_METHODOLOGY"]: - d['Inclusion of essential RBRVS "Fee Source" Language (Y/N)'] = "Y" - else: - d['Inclusion of essential RBRVS "Fee Source" Language (Y/N)'] = "N" - - # IP_OP - if "FACILITY" in d["PROV_TYPE"].upper(): - prompt = prompts.CONDITIONAL_IP_OP_SMALL( - d["FULL_SERVICE"], d["FULL_METHODOLOGY"] - ) - ip_op_answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=16 - ) - if ip_op_answer.strip() in ["IP", "OP"]: - d["IP_OP"] = ip_op_answer.strip() - else: - d["IP_OP"] = "N/A" - - if "IP" in ip_op_answer: - # IP - DSH/IME/UC, included (Y/N) - if page_num in dsh_dict.keys(): - d["IP - DSH/IME/UC, included (Y/N)"] = dsh_dict[page_num] - else: - prompt = prompts.CONDITIONAL_DSH_IME_UC(text_dict[page_num]) - dsh_answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE2, filename, max_tokens=16 - ) - d["IP - DSH/IME/UC, included (Y/N)"] = dsh_answer.strip() - dsh_dict[page_num] = dsh_answer - - # Stop-Loss Catastrophic Threshold - if page_num in sl_dict.keys(): - d["IP - Stoploss Catastrophic Threshold"] = sl_dict[page_num] - else: - if any( - [ - i.upper() in text_dict[page_num].upper() - for i in ["stop loss", "stop-loss", "stoploss"] - ] - ): - prompt = prompts.CONDITIONAL_STOP_LOSS(text_dict[page_num]) - sl_answer = llm_utils.invoke_claude( - prompt, - config.MODEL_ID_CLAUDE3_HAIKU, - filename, - max_tokens=256, - ) - else: - sl_answer = "N/A" - d["IP - Stoploss Catastrophic Threshold"] = sl_answer - sl_dict[page_num] = sl_answer - else: - d["IP - Stoploss Catastrophic Threshold"] = "N/A" - d["IP - DSH/IME/UC, included (Y/N)"] = "N/A" - else: - d["IP - Stoploss Catastrophic Threshold"] = "N/A" - d["IP - DSH/IME/UC, included (Y/N)"] = "N/A" - d["IP_OP"] = "N/A" - else: - d["IP_OP"] = "N/A" - d["IP - DSH/IME/UC, included (Y/N)"] = "N/A" - d['Inclusion of essential RBRVS "Fee Source" Language (Y/N)'] = "N/A" - d["IP - Stoploss Catastrophic Threshold"] = "N/A" - - return combined_results diff --git a/fieldExtraction/src/investment/file_processing.py b/fieldExtraction/src/investment/file_processing.py index 9048f34..7006c3a 100644 --- a/fieldExtraction/src/investment/file_processing.py +++ b/fieldExtraction/src/investment/file_processing.py @@ -3,7 +3,6 @@ import os import pandas as pd -import src.investment.conditional_funcs as conditional_funcs import src.investment.one_to_n_funcs as one_to_n_funcs import src.investment.postprocess as postprocess import src.investment.preprocess as preprocess @@ -22,21 +21,18 @@ def run_ac_prompts(file_object): print(f"Processing AC for {filename}...") ################## PREPROCESS - SMART CHUNK ################## - - # Top Sheet Addition - text_dict, top_sheet_dict, exhibit_pages, num_pages, ac_chunks = ( - preprocess.preprocess(contract_text, filename, fields="ac") - ) + contract_text = preprocess.clean_text(contract_text) + text_dict, top_sheet_dict, num_pages = preprocess.split_text(contract_text) + ac_chunks = preprocess.one_to_one_smart_chunking(text_dict, contract_text) print(f"AC Preprocessing Complete - {filename}") ################## RUN REGEX FUNCTIONALITY ################## - # no prompting for these fields - irs_answers = irs_hotfix(filename, text_dict, top_sheet_dict) ac_answers_dict.update(irs_answers) npi_answers = npi_hotfix(text_dict, top_sheet_dict) ac_answers_dict.update(npi_answers) + ################## RUN SMART CHUNKED PROMPTS ################## ac_answers_dict, no_keyword_matches = ( smart_chunking_funcs.run_smart_chunk_prompts( @@ -66,15 +62,9 @@ def run_ac_prompts(file_object): os.makedirs(output_dir, exist_ok=True) ################## POSTPROCESS ################## - ac_df = pd.DataFrame([ac_answers_dict]) - - # print(f"Before ac_postprocess Unique values: {ac_df['CONTRACT_EFFECTIVE_DT'].unique()}") - ac_df = postprocess.ac_postprocess(ac_df, text_dict, filename, num_pages) - # print(f"After ac_postprocess Unique values: {ac_df['Contract Effective Date'].unique()}") - ################## WRITE TO OUTPUT ################## ac_df.to_csv(os.path.join(output_dir, config.AC_RESULTS_NAME), index=False) @@ -105,36 +95,33 @@ def run_b_prompts(file_object): print(f"Processing B for {filename}...") ################## PREPROCESS ################## - text_dict, top_sheet_dict, exhibit_pages, num_pages, ac_chunks = ( - preprocess.preprocess(contract_text, filename, fields="b") - ) + contract_text = preprocess.clean_text(contract_text) + text_dict, top_sheet_dict, num_pages = preprocess.split_text(contract_text) + exhibit_pages, exhibit_chunk_mapping = preprocess.one_to_n_exhibit_chunking(text_dict, filename) + text_dict = preprocess.clean_tables(text_dict, filename) print(f"B Preprocessing Complete - {filename}") ################## RUN BOTTOM UP PROMPTS ################## bu_results = one_to_n_funcs.run_bottom_up( - filename, text_dict + filename, + text_dict, + exhibit_chunk_mapping ) # Returns list of dictionaries print(f"B Bottom Up Complete - {filename}") ################## RUN TOP DOWN PROMPTS ################## combined_results = one_to_n_funcs.run_top_down( - filename, text_dict, bu_results, exhibit_pages + filename, text_dict, bu_results, exhibit_chunk_mapping ) # Returns list of dictionaries print(f"B Top Down Complete - {filename}") - ################## RUN CONDITIONAL PROMPTS ################## - conditional_results = conditional_funcs.run_conditional( - combined_results, text_dict, filename - ) # List of dictionaries - print(f"B Conditional Prompts Complete - {filename}") - ################## CREATE OUTPUT DIRECTORIES ################## base_filename = os.path.splitext(filename)[0].strip() output_dir = os.path.join(config.OUTPUT_DIRECTORY, base_filename) os.makedirs(output_dir, exist_ok=True) ################## WRITE UNPROCESSED OUTPUT ################## - combined_df = pd.DataFrame(conditional_results) + combined_df = pd.DataFrame(combined_results) combined_df.to_csv( os.path.join(output_dir, config.UNPROCESSED_RESULTS_NAME), index=False ) diff --git a/fieldExtraction/src/investment/one_to_n_funcs.py b/fieldExtraction/src/investment/one_to_n_funcs.py index c82fe03..e2c8e25 100644 --- a/fieldExtraction/src/investment/one_to_n_funcs.py +++ b/fieldExtraction/src/investment/one_to_n_funcs.py @@ -9,7 +9,9 @@ from src import config, postprocessing_funcs, prompts def run_bottom_up_primary( - text_dict: dict, tokens: int, filename: str + text_dict: dict, + exhibit_chunk_mapping: dict, + filename: str ) -> dict[str, str]: """Executes the primary Bottom Up processing for pages that contain reimbursement terms. @@ -18,28 +20,25 @@ def run_bottom_up_primary( Args: text_dict (dict): A dictionary containing text content of documents keyed by page numbers. - tokens (int): The maximum number of tokens to use in the language model invocation for generating the response. Deprecated. - filename (str): The filename of the contract in question. Used for logging purposes. + exhibit_chunk_mapping (dict): A dictionary containing exhibit pages keyed by page numbers. All keys with the same exhibit page value belong to that exhibit + filename (str): The filename of the contract in question. Used for logging purposes. Returns: dict[str, str]: A dictionary where each key is a page number (string) and the value is the response from the language model. """ - - # chunk_dict = preprocess.chunk_text(text_dict) answer_strings = {} - # for page_number in chunk_dict.keys(): - # if '%' in chunk_dict[page_number] or '$' in chunk_dict[page_number]: - # prompt = prompts.BOTTOM_UP_PRIMARY(chunk_dict[page_number], config.CLIENT_NAME) - for page_number in text_dict.keys(): - if string_utils.contains_reimbursement(text_dict, page_number): + + for exhibit_page in exhibit_chunk_mapping.values(): + exhibit_chunk = string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page) + if string_utils.contains_reimbursement(exhibit_chunk): # Run Primary prompt = prompts.BOTTOM_UP_PRIMARY( - text_dict[page_number], config.CLIENT_NAME + exhibit_chunk, config.CLIENT_NAME ) answer = llm_utils.invoke_claude( prompt, config.MODEL_ID_CLAUDE35_SONNET, filename ) - answer_strings[page_number] = answer + answer_strings[exhibit_page] = answer answer_dicts = string_utils.primary_string_to_dict(answer_strings, filename) @@ -59,7 +58,11 @@ def get_short_methodology(s): def run_bottom_up_secondary( - answer_dicts: list[dict], text_dict: dict, tokens: int, filename: str + answer_dicts: list[dict], + text_dict: dict, + exhibit_chunk_mapping: dict, + tokens: int, + filename: str ) -> list[dict]: """Executes the secondary Bottom Up processing phase on the results obtained from the primary Bottom Up analysis. @@ -83,7 +86,7 @@ def run_bottom_up_secondary( final_dicts = [] for d in answer_dicts: if d is not None: - page_num = d["page_num"] + exhibit_page_num = d["page_num"] # if '%' in d['FULL_METHODOLOGY'] or '$' in d['FULL_METHODOLOGY']: prompt = prompts.BOTTOM_UP_METHODOLOGY_BREAKOUT(d) @@ -113,22 +116,15 @@ def run_bottom_up_secondary( else: d["SHORT_METHODOLOGY"] = "N/A" - # Escalator - prompt = prompts.BOTTOM_UP_ESCALATOR(d, text_dict[page_num]) - escalator_answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=1024 - ) - escalator_dict = string_utils.secondary_string_to_dict( - escalator_answer, filename - ) - d.update(escalator_dict) - final_dicts.append(d) return final_dicts -def run_bottom_up(filename: str, text_dict: dict) -> list[dict]: +def run_bottom_up(filename: str, + text_dict: dict, + exhibit_chunk_mapping: dict + ) -> list[dict]: """Processes the text of a document using a two-tiered Bottom Up approach to extract key financial and operational information. This function first runs BOTTOM_UP_PRIMARY and BOTTOM_UP_SECONDARY prompts, then processes these initial results to further @@ -137,6 +133,7 @@ def run_bottom_up(filename: str, text_dict: dict) -> list[dict]: Args: filename (str): The name of the file being processed, used to tag output data. text_dict (dict): A dictionary of text keyed by page numbers that contains the content to be analyzed. + exhibit_chunk_mapping (dict): A mapping of page_num strings (keys) to the Exhibit start page they belong to. E.g. {'1' : '1', '2' : '2', '3' : '2'}, for a 3 page contract with an exhibit start on page 2 Returns: list[dict]: A list of dictionaries with each dictionary containing refined and structured information @@ -144,15 +141,19 @@ def run_bottom_up(filename: str, text_dict: dict) -> list[dict]: """ # Bottom Up Primary - FULL_SERVICE, FULL_METHODOLOGY - bu_primary_results = run_bottom_up_primary(text_dict, 8000, filename) + bu_primary_results = run_bottom_up_primary(text_dict, + exhibit_chunk_mapping, + filename) # Bottom Up Secondary bu_secondary_results = run_bottom_up_secondary( - bu_primary_results, text_dict, 8000, filename + bu_primary_results, + text_dict, + exhibit_chunk_mapping, + 8000, + filename ) - for d in bu_secondary_results: - d["Filename"] = filename return bu_secondary_results # List of dictionaries @@ -182,188 +183,36 @@ def get_exhibit_range(text_dict, exhibit_pages, bu_page): return pages_between -def run_top_down_primary(text_dict, td_pages, filename): - page_text = "\n".join([text_dict[page_num] for page_num in td_pages]) +def run_top_down_primary(exhibit_chunk, filename): - prompt = prompts.TOP_DOWN_PRIMARY(page_text) + prompt = prompts.TOP_DOWN_PRIMARY(exhibit_chunk) answer = llm_utils.invoke_claude( prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, max_tokens=4096 ) answer_dict = string_utils.secondary_string_to_dict(answer, filename) + + # Fill unanswered fields for field in prompts.TD_PRIMARY_FIELDS: if field not in answer_dict: answer_dict[field] = "N/A" - - # Add indicators - if "ADD_ON_REIMBURSEMENT_LANGUAGE" in answer_dict: - if "N/A" in answer_dict["ADD_ON_REIMBURSEMENT_LANGUAGE"]: - answer_dict["ADD_ON_REIMBURSEMENT_IND"] = "N" - else: - answer_dict["ADD_ON_REIMBURSEMENT_IND"] = "Y" - else: - answer_dict["ADD_ON_REIMBURSEMENT_IND"] = "N" - answer_dict["ADD_ON_REIMBURSEMENT_LANGUAGE"] = "N/A" - + + # Fill Prov Type if specified by client if config.PROV_TYPE: answer_dict["PROV_TYPE"] = config.PROV_TYPE return answer_dict -def get_td_dict(text_dict, filename, PROMPT_FUNC, VALID_VALUES): - td_dict = {} - for page_num, page_text in text_dict.items(): - if ( - any(i.upper() in page_text.upper() for i in VALID_VALUES) - and page_num.isdigit() - ): - prompt = PROMPT_FUNC(page_text) - answer = llm_utils.invoke_claude( - prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 1024 - ) - td_dict[page_num] = answer - return td_dict - - -def add_medically_necessary(merged_dicts, medically_necessary_dict): - if not medically_necessary_dict: - for d in merged_dicts: - d["MEDICAL_NECESSITY_LANGUAGE"] = "N/A" - d["MEDICAL_NECESSITY_LANGUAGE_IND"] = "N" - return merged_dicts - - sorted_keys = sorted(int(k) for k in medically_necessary_dict.keys()) - - # Find the closest page number that is less than or equal to the given page number - def find_closest_page(page_num): - closest_page = sorted_keys[0] - for key in sorted_keys: - if key <= page_num: - closest_page = key - else: - break - return closest_page - - for d in merged_dicts: - try: - current_page_num = int(d["page_num"]) - # Find the closest preceding medically necessary page - closest_page = find_closest_page(current_page_num) - d["MEDICAL_NECESSITY_LANGUAGE"] = medically_necessary_dict[ - str(closest_page) - ] - if not string_utils.is_empty(d["MEDICAL_NECESSITY_LANGUAGE"]): - d["MEDICAL_NECESSITY_LANGUAGE_IND"] = "Y" - else: - d["MEDICAL_NECESSITY_LANGUAGE_IND"] = "N" - except: - d["MEDICAL_NECESSITY_LANGUAGE"] = "N/A" - d["MEDICAL_NECESSITY_LANGUAGE_IND"] = "N" - - return merged_dicts - - -def add_exclusions(merged_dicts, exclusions_dict): - if not exclusions_dict: - for d in merged_dicts: - d["EXCLUSIONS"] = "N/A" - return merged_dicts - - # Prepare by sorting exclusion keys and initializing an empty dictionary to track exclusions by EXHIBIT - sorted_keys = sorted(int(k) for k in exclusions_dict.keys()) - exhibit_exclusions = {} - - # First pass: assign exclusions based on page number - for d in merged_dicts: - page_num = str(d["page_num"]) - if page_num.isdigit(): - for exc_page in sorted_keys: - if exc_page >= int(page_num) and exc_page <= int(page_num) + 2: - new_exclusion = exclusions_dict[str(exc_page)] - # Handle EXCLUSIONS as a list to avoid duplicates - if "EXCLUSIONS" not in d: - d["EXCLUSIONS"] = [] - if new_exclusion not in d["EXCLUSIONS"]: - try: - d["EXCLUSIONS"].append(new_exclusion) - except: - d["EXCLUSIONS"] = [new_exclusion] - - # Map this exclusion to the EXHIBIT value - if "EXHIBIT" in d: - if d["EXHIBIT"] not in exhibit_exclusions: - exhibit_exclusions[d["EXHIBIT"]] = [] - if new_exclusion not in exhibit_exclusions[d["EXHIBIT"]]: - exhibit_exclusions[d["EXHIBIT"]].append(new_exclusion) - break - else: - # If no matching exclusion found, check if 'EXHIBIT' maps to any existing exclusion - if "EXHIBIT" in d and d["EXHIBIT"] in exhibit_exclusions: - for exclusion in exhibit_exclusions[d["EXHIBIT"]]: - if "EXCLUSIONS" not in d: - d["EXCLUSIONS"] = [] - if exclusion not in d["EXCLUSIONS"]: - d["EXCLUSIONS"].append(exclusion) - else: - d["EXCLUSIONS"] = ["N/A"] - else: - d["EXCLUSIONS"] = ["N/A"] - - # Convert lists to semicolon-separated strings - for d in merged_dicts: - if isinstance(d["EXCLUSIONS"], list): - d["EXCLUSIONS"] = "; ".join(d["EXCLUSIONS"]) - - return merged_dicts - - -def run_top_down(filename, text_dict, bu_results, exhibit_pages): - bu_pages = {bu_dict["page_num"] for bu_dict in bu_results} - - td_page_dict = {} +def run_top_down(filename, text_dict, bu_results, exhibit_chunk_mapping): td_answer_dict = {} - - for bu_page in bu_pages: - td_pages = get_exhibit_range(text_dict, exhibit_pages, bu_page) - if td_pages not in td_page_dict.values(): - answer_dict = run_top_down_primary(text_dict, td_pages, filename) - td_page_dict[bu_page] = td_pages - td_answer_dict[bu_page] = answer_dict - else: - td_page_dict[bu_page] = td_pages - for key, value in td_page_dict.items(): - if value == td_pages: - answer_dict = td_answer_dict[key] - td_answer_dict[bu_page] = answer_dict - break - merged_dicts = [] for bu_dict in bu_results: - answer_dict = td_answer_dict[bu_dict["page_num"]] + bu_exhibit_page = bu_dict['page_num'] + if bu_exhibit_page in td_answer_dict.keys(): # When the exhibit has already had top down prompts run + answer_dict = td_answer_dict[bu_exhibit_page] + else: # When the exhibit has not yet had top down prompts run + exhibit_chunk = string_utils.get_exhibit_chunk(text_dict, exhibit_chunk_mapping, bu_exhibit_page) + answer_dict = run_top_down_primary(exhibit_chunk, filename) + td_answer_dict[bu_exhibit_page] = answer_dict bu_dict.update(answer_dict) merged_dicts.append(bu_dict) - - # Run and Merge Exclusions - exclusion_dict = get_td_dict( - text_dict, filename, prompts.TOP_DOWN_EXCLUSIONS, valid.EXCLUSION_TERMS - ) - final_dicts = add_exclusions(merged_dicts, exclusion_dict) - - # Run and Merge Medically Necessary - medically_necessary_dict = get_td_dict( - text_dict, - filename, - prompts.TOP_DOWN_MEDICALLY_NECESSARY, - valid.VALID_MEDICALLY_NECESSARY, - ) - medically_necessary_dict = postprocessing_funcs.filter_dict( - medically_necessary_dict, - pattern=re.compile( - "|".join( - re.escape(keyword) for keyword in valid.INVALID_MEDICALLY_NECESSARY - ), - re.IGNORECASE, - ), - ) - final_dicts = add_medically_necessary(final_dicts, medically_necessary_dict) - - return final_dicts + return merged_dicts diff --git a/fieldExtraction/src/investment/postprocess.py b/fieldExtraction/src/investment/postprocess.py index ca72571..c9bb04a 100644 --- a/fieldExtraction/src/investment/postprocess.py +++ b/fieldExtraction/src/investment/postprocess.py @@ -6,46 +6,12 @@ def b_postprocess(filename, df, pages): if df.shape[0] > 0: # Metadata fields - df.drop("Filename", axis=1, inplace=True) - df["Contract Name"] = "Filename: " + filename - df["Parent Agreement Code"] = postprocessing_funcs.get_parent_agreement_code( - filename - ) df["Pages"] = pages - # Add Single Code, Multiple Rate - df = postprocessing_funcs.add_scmr(df) + # Removed postprocessing for Client columns - # Filter Add Ons - df = postprocessing_funcs.filter_add_ons(df) - - # Clean MSR - df = postprocessing_funcs.clean_msr_lesser(df) - - # Clean Default - df = postprocessing_funcs.clean_default_term(df) - - # Clean Prov 2 - df = postprocessing_funcs.clean_prov_2(df) - - # Clean Lesser Of Rate - df = postprocessing_funcs.clean_lesser_rate(df) - - # Clean LOB - df = postprocessing_funcs.clean_lob(df, filename) - - df = postprocessing_funcs.methodology_short_fix(df) - - df = postprocessing_funcs.yn_fixes(df) - - # Rename and reorder - df.rename(columns=valid.B_MAPPING, inplace=True) - - column_order = [col for col in valid.B_MAPPING.values() if col in df.columns] - final_df = df[column_order] - - return final_df + return df else: return df diff --git a/fieldExtraction/src/investment/preprocess.py b/fieldExtraction/src/investment/preprocess.py index 924a135..a861404 100644 --- a/fieldExtraction/src/investment/preprocess.py +++ b/fieldExtraction/src/investment/preprocess.py @@ -4,34 +4,97 @@ import src.utils.table_utils as table_utils from src import config, keywords, preprocessing_funcs -def preprocess(contract_text, filename, fields=config.FIELDS): - +def clean_text(contract_text): contract_text = preprocessing_funcs.clean_newlines(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 = {k: v for k, v in text_dict.items() if k.isdigit() and 42 <= int(k) <= 42} num_pages = len(text_dict.keys()) text_dict, top_sheet_dict = preprocessing_funcs.filter_quick_review(text_dict) + return text_dict, top_sheet_dict, num_pages - if fields == "b": - exhibit_pages = preprocessing_funcs.get_exhibit_pages( + +def one_to_n_exhibit_chunking(text_dict, filename): + """ + Identifies exhibit pages in a document and maps all text pages to their respective exhibits. + + This function identifies pages that contain exhibit headers from a given document and then + maps each page of the document to the nearest preceding exhibit header. + + Parameters: + text_dict (dict): A dictionary where keys are page numbers and values are the text on those pages. + filename (str): The name of the file being processed + + Returns: + tuple: A tuple containing two elements: + - list: A list of page numbers that are identified as starting new exhibits. + - dict: A dictionary mapping each page number to the exhibit it belongs to, with exhibit pages as keys. + + Notes: + - This function utilizes `preprocessing_funcs.get_exhibit_pages` to identify pages that start new exhibits based on the document text and filename. + - It uses `preprocessing_funcs.chunk_by_exhibit` to map each page to its respective exhibit based on the identified exhibit pages. + """ + exhibit_pages = preprocessing_funcs.get_exhibit_pages( text_dict, filename ) # All pages with exhibit headers - text_dict = table_utils.align_and_format_tables(text_dict, filename) - text_dict = preprocessing_funcs.chunk_consecutive(text_dict, exhibit_pages) - # write text_dict to a csv file: - with open("text_dict.csv", "w") as f: - writer = csv.writer(f) - for key, value in text_dict.items(): - writer.writerow([key, value]) + exhibit_chunk_mapping = preprocessing_funcs.chunk_by_exhibit(text_dict, exhibit_pages) + return exhibit_pages, exhibit_chunk_mapping - ac_chunks, full_context = "", "" - if fields == "ac": - exhibit_pages = [] - ac_chunks = preprocessing_funcs.smart_chunk_ac(text_dict, contract_text, keywords.GROUPED_KEYWORD_MAPPINGS) - return text_dict, top_sheet_dict, exhibit_pages, num_pages, ac_chunks +def clean_tables(text_dict, filename): + """ + In current state, this function is a wrapper for align_and_format_tables. Call any new table-related preprocessing funcs here + + Parameters: + text_dict (dict): A dictionary where keys are page numbers and values are the text on those pages. + filename (str): The name of the file being processed + + Returns: + dict: text_dict, with the tables aligned and formatted + """ + text_dict = table_utils.align_and_format_tables(text_dict, filename) + return text_dict + + +def one_to_one_smart_chunking(text_dict, contract_text, keyword_mappings=keywords.GROUPED_KEYWORD_MAPPINGS): + """ + In current state, this function is a wrapper for smart_chunk_ac. Call any new smart-chunking-related preprocessing funcs here + + Parameters: + text_dict (dict): A dictionary where keys are page numbers and values are the text on those pages. + contract_text (str): The full text of the contract + keyword_mappings (dict): keywords.GROUPED_KEYWORD_MAPPINGS by default + + Returns: + dict: ac_chunks + """ + ac_chunks = preprocessing_funcs.smart_chunk_ac(text_dict, contract_text, keyword_mappings) + return ac_chunks + diff --git a/fieldExtraction/src/preprocessing_funcs.py b/fieldExtraction/src/preprocessing_funcs.py index 1cb21d2..a115e30 100644 --- a/fieldExtraction/src/preprocessing_funcs.py +++ b/fieldExtraction/src/preprocessing_funcs.py @@ -307,3 +307,39 @@ def get_exhibit_pages(text_dict, filename): 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_dict = {} + current_exhibit = "0" + for page_num in text_dict.keys(): + if page_num in exhibit_pages: + current_exhibit = page_num + exhibit_dict[page_num] = current_exhibit + else: + exhibit_dict[page_num] = current_exhibit + return exhibit_dict + + diff --git a/fieldExtraction/src/utils/string_utils.py b/fieldExtraction/src/utils/string_utils.py index d32712c..269f51c 100644 --- a/fieldExtraction/src/utils/string_utils.py +++ b/fieldExtraction/src/utils/string_utils.py @@ -219,7 +219,6 @@ def primary_string_to_dict(string_dict, filename): dict_list = [secondary_string_to_dict(d, filename) for d in dict_list] for dict_ in dict_list: dict_["page_num"] = page_num - dict_["Filename"] = filename data.append(dict_) return data @@ -250,4 +249,30 @@ def is_empty(value): # string_funcs.py return True else: empty_values = [None, "", "N/A", "NA", "null", "none", "NaN", np.nan, "nan"] - return value in empty_values \ No newline at end of file + return value in empty_values + + +def get_exhibit_chunk(text_dict: dict, + exhibit_chunk_mapping: dict, + exhibit_page: str) -> str: + """ + Compiles text from a specific exhibit based on its starting page number. + + This function aggregates the text of all pages that are mapped to a specific exhibit + start page in a document. It uses the exhibit chunk mapping to determine which pages + belong to the exhibit defined by the given start page number, and concatenates their + texts into a single string. + + Parameters: + text_dict (dict): A dictionary where keys are page numbers and values are the text on those pages. + exhibit_chunk_mapping (dict): A dictionary mapping each page number to the starting page number of the exhibit it belongs to. + exhibit_page (str): The page number that marks the beginning of the exhibit to compile text for. + + Returns: + str: A string that combines all the text from the pages belonging to the specified exhibit, separated by newlines. + + Notes: + - The function assumes that `exhibit_chunk_mapping` correctly maps all page numbers in `text_dict` to their respective exhibit starting pages. + - It is important that `exhibit_page` exists as a key in the `exhibit_chunk_mapping` dictionary and corresponds to the starting page of an exhibit. + """ + return '\n'.join([page_text for page_num, page_text in text_dict.items() if exhibit_chunk_mapping[page_num] == exhibit_page]) diff --git a/fieldExtraction/tests/exhibit_chunk_mapping_test.py b/fieldExtraction/tests/exhibit_chunk_mapping_test.py new file mode 100644 index 0000000..f8d4da5 --- /dev/null +++ b/fieldExtraction/tests/exhibit_chunk_mapping_test.py @@ -0,0 +1,31 @@ + +from src.preprocessing_funcs import chunk_by_exhibit + +def test1(): + """Test non-empty text_dict with empty exhibit_pages.""" + text_dict = {'1': 'This', '2': 'is', '3': 'an', '4': 'exhibit', '5': 'test'} + exhibit_pages = [] + assert chunk_by_exhibit(text_dict, exhibit_pages) == {'1': '1', '2': '2', '3': '3', '4': '4', '5': '5'} + +def test2(): + """Test when exhibit_pages contains some keys that are also in text_dict.""" + text_dict = {'1': 'Welcome', '2': 'to', '3': 'the', '4': 'exhibit'} + exhibit_pages = ['1', '2', '4'] + expected_output = {'1': '1', '2': '2', '3': '2', '4': '4'} + assert chunk_by_exhibit(text_dict, exhibit_pages) == expected_output + + +def test3(): + """Test multiple exhibit pages with transitions between them.""" + text_dict = {'1': 'Page 1', '2': 'Page 2', '3': 'Page 3', '4': 'Page 4', '5': 'Page 5'} + exhibit_pages = ['2', '4'] + expected_output = {'1': '0', '2': '2', '3': '2', '4': '4', '5': '4'} + assert chunk_by_exhibit(text_dict, exhibit_pages) == expected_output + + +test1() + +test2() + +test3() + diff --git a/fieldExtraction/tests/get_exhibit_chunk_test.py b/fieldExtraction/tests/get_exhibit_chunk_test.py new file mode 100644 index 0000000..86a7c44 --- /dev/null +++ b/fieldExtraction/tests/get_exhibit_chunk_test.py @@ -0,0 +1,34 @@ + +from src.utils.string_utils import get_exhibit_chunk + +def test1(): + """Test non-empty text_dict with empty exhibit_pages.""" + text_dict = {'1': 'This', '2': 'is', '3': 'an', '4': 'exhibit', '5': 'test'} + exhibit_chunk_mapping = {'1': '1', '2': '2', '3': '3', '4': '4', '5': '5'} + exhibit_page_num = '4' + assert get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page_num) == 'exhibit' + +def test2(): + """Test when exhibit_pages contains some keys that are also in text_dict.""" + text_dict = {'1': 'Welcome', '2': 'to', '3': 'the', '4': 'exhibit'} + exhibit_chunk_mapping = {'1': '1', '2': '2', '3': '2', '4': '4'} + exhibit_page_num = '2' + expected_output = "to\nthe" + assert get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page_num) == expected_output + + +def test3(): + """Test multiple exhibit pages with transitions between them.""" + text_dict = {'1': 'Page 1', '2': 'Page 2', '3': 'Page 3', '4': 'Page 4', '5': 'Page 5'} + exhibit_chunk_mapping = {'1': '0', '2': '2', '3': '2', '4': '4', '5': '4'} + exhibit_page_num = '4' + expected_output = 'Page 4\nPage 5' + assert get_exhibit_chunk(text_dict, exhibit_chunk_mapping, exhibit_page_num) == expected_output + + +test1() + +test2() + +test3() + diff --git a/fieldExtraction/tests/string_utils_test.py b/fieldExtraction/tests/string_utils_test.py index 2f7fd77..6effc02 100644 --- a/fieldExtraction/tests/string_utils_test.py +++ b/fieldExtraction/tests/string_utils_test.py @@ -96,9 +96,9 @@ class TestPrimaryStringToDict: } filename = "test_file.txt" expected_output = [ - {"key1": "value1", "page_num": "1", "Filename": "test_file.txt"}, - {"key2": "value2", "page_num": "1", "Filename": "test_file.txt"}, - {"key3": "value3", "page_num": "2", "Filename": "test_file.txt"}, + {"key1": "value1", "page_num": "1"}, + {"key2": "value2", "page_num": "1"}, + {"key3": "value3", "page_num": "2"}, ] result = primary_string_to_dict(string_dict, filename) assert result == expected_output @@ -135,8 +135,8 @@ class TestPrimaryStringToDict: } filename = "test_file.txt" expected_output = [ - {"key1": "value1", "page_num": "1", "Filename": "test_file.txt"}, - {"key2": "value2", "page_num": "2", "Filename": "test_file.txt"}, + {"key1": "value1", "page_num": "1"}, + {"key2": "value2", "page_num": "2"}, ] result = primary_string_to_dict(string_dict, filename) assert result == expected_output @@ -155,10 +155,10 @@ class TestPrimaryStringToDict: } filename = "test_file.txt" expected_output = [ - {"key1": "value1", "page_num": "1", "Filename": "test_file.txt"}, - {"key2": "value2", "page_num": "1", "Filename": "test_file.txt"}, - {"key3": "value3", "page_num": "2", "Filename": "test_file.txt"}, - {"key4": "value4", "page_num": "2", "Filename": "test_file.txt"}, + {"key1": "value1", "page_num": "1"}, + {"key2": "value2", "page_num": "1"}, + {"key3": "value3", "page_num": "2"}, + {"key4": "value4", "page_num": "2"}, ] result = primary_string_to_dict(string_dict, filename) assert result == expected_output