doczyai-pipelines/fieldExtraction/src/investment/file_processing.py

import csv
import os

import pandas as pd

import src.investment.one_to_n_funcs as one_to_n_funcs
import src.investment.one_to_one_funcs as one_to_one_funcs
import src.investment.postprocess as postprocess
import src.investment.preprocess as preprocess
import src.investment.aarete_derived as aarete_derived
import src.investment.smart_chunking_funcs as smart_chunking_funcs
from src.prompts.investment_prompts import Field, FieldSet
from src import config
from src.regex.regex_utils import irs_hotfix, npi_hotfix
import src.utils.string_utils as string_utils
import src.utils.io_utils as io_utils
from src.config import FIELD_JSON_PATH


def process_file(file_object, run_timestamp):
    filename, contract_text = file_object
    print(f"Processing {filename}...")

    ################## PREPROCESS ##################
    contract_text = preprocess.clean_text(contract_text)
    text_dict, top_sheet_dict = preprocess.split_text(contract_text)
    text_dict = preprocess.clean_tables(text_dict, filename) # TODO: Optimize later to not run a bunch of extra prompts for the extra table pages
    exhibit_pages, exhibit_chunk_mapping = preprocess.one_to_n_exhibit_chunking(text_dict, filename)
    print(f"Preprocessing Complete - {filename}")

    ################## ONE TO ONE ##################
    one_to_one_results = run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict) # Return df
    one_to_one_results['CONTRACT_FILE_NAME'] = filename

    ################## ONE TO N ##################
    if string_utils.contains_reimbursement(contract_text):
        one_to_n_results = run_one_to_n_prompts(filename, text_dict, exhibit_pages, exhibit_chunk_mapping) # Return df
        one_to_n_results['CONTRACT_FILE_NAME'] = filename
        contract_results = pd.merge(one_to_one_results, one_to_n_results, how='right', on='CONTRACT_FILE_NAME')
    else:
        contract_results = one_to_n_results

    ################## AARETE-DERIVED ##################
    final_results = aarete_derived.get_aarete_derived(contract_results)
    
    ################## POSTPROCESS ##################
    final_df = postprocess.postprocess(final_results)

    ################## WRITE INDIVIDUAL ##################
    if config.WRITE_TO_S3:
        io_utils.write_s3(final_df, filename, run_timestamp, "individual")
    else:
        io_utils.write_local(final_df, filename, "", "individual")

    return final_df


def run_one_to_one_prompts(filename, contract_text, text_dict, top_sheet_dict):

    ################## INITIALIZE FIELDS ##################
    one_to_one_fields = FieldSet(relationship="one_to_one", file_path=config.FIELD_JSON_PATH)

    ################## RUN REGEX ##################
    regex_answers_dict = one_to_one_funcs.run_regex_fields(one_to_one_fields, contract_text, filename)
    
    ################## RUN SMART CHUNKED PROMPTS ##################
    smart_chunked_answers_dict = one_to_one_funcs.run_smart_chunked_fields(
        one_to_one_fields, contract_text, filename
    )

    ################## RUN FULL CONTEXT PROMPTS ##################
    full_context_answers_dict = one_to_one_funcs.run_full_context_fields(
        one_to_one_fields, contract_text, filename
    )

    ################## CONVERT TO DF ##################
    final_answers_dict = {**regex_answers_dict, **smart_chunked_answers_dict, **full_context_answers_dict}
    final_df = pd.DataFrame([final_answers_dict])

    return final_df


def run_one_to_n_prompts(filename, text_dict, exhibit_pages, exhibit_chunk_mapping):

    ################## RUN PROMPTS ##################
    one_to_n_results = one_to_n_funcs.get_one_to_n(text_dict, filename, exhibit_pages, exhibit_chunk_mapping)

    one_to_n_df = pd.DataFrame(one_to_n_results)
    return one_to_n_df