doczyai-pipelines/fieldExtraction/src/file_processing.py

import csv
import os

import pandas as pd

import conditional_funcs
import config
import one_to_n_funcs
import postprocess
import preprocess
import smart_chunking_funcs
from regex_funcs import irs_hotfix, npi_hotfix


def run_ac_prompts(file_object):

    ac_answers_dict = {}
    log_data = []
    chunk_data = []

    ################## INITIATE PROCESSING ##################
    filename, contract_text = 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")
    )
    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(
            ac_answers_dict, ac_chunks, contract_text, filename, text_dict
        )
    )

    ################## GET FULL CONTEXT FIELDS ##################
    # Process 'full_context' fields together
    full_context_fields = smart_chunking_funcs.get_full_context_fields(
        ac_answers_dict, no_keyword_matches
    )

    ################## RUN FULL CONTEXT PROMPTS ##################
    ac_answers_dict = smart_chunking_funcs.run_full_context_prompts(
        ac_answers_dict, full_context_fields, contract_text, filename
    )

    ################## AC CONDITIONAL PROMPTS ##################
    ac_answers_dict = smart_chunking_funcs.run_ac_conditional_prompts(
        ac_answers_dict, contract_text, 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)

    ################## 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)

    chunk_file_path = "chunk_log.csv"
    with open(chunk_file_path, "a", newline="", encoding="utf-8") as chunk_file:
        fieldnames = [
            "Contract name",
            "Field group",
            "Field list",
            "Methodology",
            "Case sensitivity",
            "Page list",
            "Chunk size",
            "% Reduction",
        ]
        writer = csv.DictWriter(chunk_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(chunk_data)

    print(f"Chunk log written to {chunk_file_path}")

    return ac_answers_dict


def run_b_prompts(file_object):
    ################## INITIATE PROCESSING ##################
    filename, contract_text = 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")
    )
    print(f"B Preprocessing Complete - {filename}")

    ################## RUN BOTTOM UP PROMPTS ##################
    bu_results = one_to_n_funcs.run_bottom_up(
        filename, text_dict
    )  # 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
    )  # 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.to_csv(
        os.path.join(output_dir, config.UNPROCESSED_RESULTS_NAME), index=False
    )

    ################## RUN POSTPROCESSING ##################
    final_df = postprocess.b_postprocess(filename, combined_df, num_pages)
    print(f"B Postprocessing Complete - {filename} ")

    ################## WRITE FINAL ##################
    final_df.to_csv(os.path.join(output_dir, config.B_RESULTS_NAME), index=False)
    print(f"B Output Complete - {filename} ")