Katon Minhas
137 lines
5.3 KiB
Python
137 lines
5.3 KiB
Python
|
|
import os
|
|
import pandas as pd
|
|
|
|
import config
|
|
import preprocess
|
|
import table_funcs
|
|
import prompt_funcs
|
|
import prompts
|
|
import claude_funcs
|
|
import postprocessingfuncs
|
|
import postprocess
|
|
|
|
def clean_td(td):
|
|
td_clean = []
|
|
for d in td:
|
|
new_d = {}
|
|
for k, v in d.items():
|
|
if 'DATE' in k:
|
|
new_d[k] = v if isinstance(v, list) else [v]
|
|
elif k not in ['page_num', 'Filename']:
|
|
if isinstance(v, str) and ',' in v:
|
|
new_d[k] = [item.strip() for item in v.split(',')]
|
|
elif v == 'N/A':
|
|
new_d[k] = []
|
|
else:
|
|
new_d[k] = [v] if isinstance(v, str) else v
|
|
else:
|
|
new_d[k] = v
|
|
td_clean.append(new_d)
|
|
return td_clean
|
|
|
|
|
|
def get_unique_keys(dicts):
|
|
keys = set()
|
|
for d in dicts:
|
|
keys.update(d.keys())
|
|
return keys
|
|
|
|
def get_unique_date_fields(td_results):
|
|
unique_date_fields = {}
|
|
for td in td_results:
|
|
for key, value in td.items():
|
|
if 'DATE' in key:
|
|
if key not in unique_date_fields:
|
|
unique_date_fields[key] = set()
|
|
unique_date_fields[key].update(value if isinstance(value, list) else [value])
|
|
for key in unique_date_fields:
|
|
unique_date_fields[key] = list(unique_date_fields[key])
|
|
return unique_date_fields
|
|
|
|
def merge_results(td_results, bu_results, text_dict):
|
|
all_keys = get_unique_keys(td_results) | get_unique_keys(bu_results)
|
|
|
|
date_fields = get_unique_date_fields(td_results)
|
|
|
|
merged_results = []
|
|
for bu in bu_results:
|
|
merged = {key: bu.get(key, "") for key in all_keys}
|
|
td_on_page = [td for td in td_results if td['page_num'] == bu['page_num']]
|
|
|
|
for td in td_on_page:
|
|
for key, value in td.items():
|
|
if not merged[key]:
|
|
merged[key] = value
|
|
elif isinstance(value, list) and value and not isinstance(merged[key], list):
|
|
merged[key] = value
|
|
elif isinstance(value, list) and value:
|
|
merged[key].extend(value)
|
|
|
|
for date_key, date_values in date_fields.items():
|
|
if date_key not in merged or not merged[date_key]:
|
|
merged[date_key] = date_values
|
|
else:
|
|
merged[date_key].extend([val for val in date_values if val not in merged[date_key]])
|
|
|
|
for key in all_keys:
|
|
if isinstance(merged[key], list) and len(merged[key]) > 1:
|
|
page_num = bu['page_num']
|
|
page_text = text_dict.get(page_num, "")
|
|
prompt = prompts.GENERATE_PROMPT(bu, td_on_page, page_text, field_name=key, values=merged[key])
|
|
response = claude_funcs.invoke_claude_3(prompt, max_tokens=4000)
|
|
try:
|
|
selected_value = response.strip()
|
|
merged[key] = selected_value
|
|
except Exception as e:
|
|
print(f"Error processing LLM response for key {key}: {e}")
|
|
merged[key] = ', '.join(merged[key])
|
|
|
|
merged_results.append(merged)
|
|
|
|
return merged_results
|
|
|
|
|
|
def process_file(file_object):
|
|
|
|
################## INITIATE PROCESSING ##################
|
|
filename, contract_text = file_object
|
|
if config.VERBOSE: print(f"Processing {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)
|
|
|
|
################## PREPROCESS ##################
|
|
contract_text = preprocess.clean_newlines(contract_text)
|
|
text_dict = preprocess.split_text(contract_text)
|
|
text_dict = table_funcs.align_and_format_tables(text_dict)
|
|
text_dict = preprocess.highlight_rates(text_dict)
|
|
|
|
################## RUN TOP DOWN ##################
|
|
td_results = prompt_funcs.run_top_down(filename, text_dict) # Returns list of dictionaries for each page
|
|
pd.DataFrame(td_results).to_csv(os.path.join(output_dir, config.TD_RESULTS_NAME), index=False)
|
|
print(f"Top Down Complete - {filename}")
|
|
|
|
################## RUN BOTTOM UP ##################
|
|
bu_results = prompt_funcs.run_bottom_up(filename, text_dict) # Returns list of dictionaries
|
|
pd.DataFrame(bu_results).to_csv(os.path.join(output_dir, config.BU_RESULTS_NAME), index=False)
|
|
print(f"Bottom Up Complete - {filename}")
|
|
|
|
################## COMBINE TOP DOWN AND BOTTOM UP ##################
|
|
combined_results = merge_results(clean_td(td_results), bu_results, text_dict)
|
|
print(f"TD/BU Merge Complete - {filename} ")
|
|
|
|
################## WRITE UNPROCESSED OUTPUT ##################
|
|
combined_df = pd.DataFrame(combined_results)
|
|
combined_df.to_csv(os.path.join(output_dir, config.UNPROCESSED_RESULTS_NAME), index=False)
|
|
|
|
################## RUN POSTPROCESSING ##################
|
|
combined_df = combined_df.applymap(postprocessingfuncs.sanitize_value)
|
|
post_processed_combined_df = postprocess.postprocess_results(combined_df)
|
|
post_processed_combined_df.to_csv(os.path.join(output_dir, config.PROCESSED_RESULTS_NAME), index=False)
|
|
print(f"Postprocessing Complete - {filename} ")
|
|
|
|
print(f"Output Complete - {filename} ")
|
|
|