doczyai-pipelines/fieldExtraction/src/utils.py

import os
import re
import pandas as pd
import shutil

import config


def read_local(file_path):
    # Check if the file is a text file
    if os.path.isfile(file_path) and file_path.endswith(".txt"):
        try:
            # First attempt to open the file with UTF-8 encoding
            with open(file_path, "r", encoding="utf-8") as file:
                file_contents = file.read()
                return file_contents
        except UnicodeDecodeError:
            # If UTF-8 fails, try reading the file with ANSI encoding
            try:
                with open(file_path, "r", encoding="cp1252") as file:
                    file_contents = file.read()
                    return file_contents
            except UnicodeDecodeError:
                # If ANSI also fails, log an error message or handle it accordingly
                print(f"Failed to decode {file_path} with UTF-8 and cp1252 encodings.")


def read_s3():
    s3_client = config.S3_CLIENT
    objects = s3_client.list_objects_v2(Bucket=config.BUCKET, Prefix=config.PREFIX)
    file_list = []
    for obj in objects["Contents"]:
        if not obj["Key"].endswith("/"):
            file_list.append(obj["Key"])
    contract_list = sorted(file_list)
    files = {}
    for contract in contract_list:
        data = s3_client.get_object(Bucket=config.BUCKET, Key=contract)
        contents = data["Body"].read()
        context = contents.decode("utf-8")
        path, filename = os.path.split(contract)
        files[filename] = context
    return files


def read_input(path=config.LOCAL_PATH, mode=config.READ_MODE):
    if mode == "_LOCAL_":
        files = {}
        for file in os.listdir(path):
            full_path = os.path.join(path, file)
            file_text = read_local(full_path)
            files[file] = file_text
        return files
    elif mode == "_S3_":
        return read_s3()


def consolidate_individual(input_folder="results", output_folder="output"):
    dfs = []
    for filename in os.listdir(input_folder):
        if filename.endswith(".csv"):
            filepath = os.path.join(input_folder, filename)
            dfs.append(pd.read_csv(filepath))

    # Remove temp folder
    if input_folder == "temp" and os.path.exists(input_folder):
        shutil.rmtree(input_folder)

    consolidated_df = pd.concat(dfs, ignore_index=True)

    # Write output
    version = 1
    existing_files = [
        filename
        for filename in os.listdir(output_folder)
        if filename.startswith(f"consolidated_results_{config.TODAY}")
    ]
    if existing_files:
        versions = [
            int(file.split("_v")[1].split(".")[0])
            for file in existing_files
            if "_v" in file
        ]
        if versions:
            version = max(versions) + 1
    filename = f"consolidated_results_{config.TODAY}_v{version}.csv"
    consolidated_df.to_csv(os.path.join(output_folder, filename))


def preprocess_text_file(file_path):
    # Attempt to open the file with UTF-8 encoding first
    try:
        with open(file_path, "r", encoding="utf-8") as file:
            text = file.read()
    except UnicodeDecodeError:
        # If UTF-8 fails, try reading the file with ANSI (cp1252) encoding
        try:
            with open(file_path, "r", encoding="cp1252") as file:
                text = file.read()
        except UnicodeDecodeError:
            # If ANSI also fails, log an error message or handle it accordingly
            print(f"Failed to decode {file_path} with UTF-8 and cp1252 encodings.")
            return []

    # Split the text into pages based on a specific marker
    pages = re.split(r"Start of Page No\. = \d+", text)
    return pages


def format_td_check(td_dicts, dont_include_list):
    final_str = ""
    dict_count = 1
    for td_dict in td_dicts:
        final_str += str(dict_count) + ". "
        for k in td_dict.keys():
            if k not in dont_include_list:
                final_str += k + ": " + td_dict[k] + ", "
        final_str += "\n"
        dict_count += 1
    return final_str


def consolidate_csvs(
    output_dir=config.CONSOLIDATED_OUTPUT_DIRECTORY, output_file=config.OUTPUT_CSV_PATH
):
    os.makedirs(output_dir, exist_ok=True)

    df_list = []
    # Walk through each folder in the output directory
    for root, dirs, files in os.walk(output_dir):
        for dir_name in dirs:
            dir_path = os.path.join(root, dir_name)
            file_path = os.path.join(dir_path, config.PROCESSED_RESULTS_NAME)
            if os.path.isfile(file_path):
                try:
                    df = pd.read_csv(file_path)
                    df_list.append(df)
                except:
                    pass

    concatenated_df = pd.concat(df_list, ignore_index=True)
    concatenated_df.to_csv(
        os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, output_file), index=False
    )
    print(f"All CSV files have been consolidated into {output_file}")


def contains_reimbursement(text, page):
    if isinstance(text, dict):
        return page.isdigit() and (
            "%" in text[page] or "$" in text[page] or "percent" in text[page].lower()
        )
    elif isinstance(text, str):
        return "%" in text or "$" in text or "percent" in text.lower()
    else:
        print("contains_reimbursement - Invalid data type")


def filter_already_processed(input_dict):
    already_processed = []
    for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
        if config.PROCESSED_RESULTS_NAME in os.listdir(
            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
        ):
            already_processed.append(folder_name + ".txt")

    input_dict = {
        key: input_dict[key]
        for key in input_dict.keys()
        if key not in already_processed
    }
    return input_dict