doczyai-pipelines/fieldExtraction/src/utils.py

import os
import re
import pandas as pd
import shutil
import time
import sys
import numpy as np
from collections import defaultdict
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 list_s3_files(prefix=config.S3_PREFIX, bucket=config.S3_BUCKET):
    """List all files in an S3 folder."""
    response = config.S3_CLIENT.list_objects_v2(Bucket=bucket, Prefix=prefix)
    return [item["Key"] for item in response.get("Contents", [])]


def read_s3():
    s3_client = config.S3_CLIENT
    bucket = config.S3_BUCKET
    prefix = config.S3_PREFIX

    file_list = []
    continuation_token = None

    while True:

        if continuation_token:
            response = s3_client.list_objects_v2(
                Bucket=bucket, Prefix=prefix, ContinuationToken=continuation_token
            )
        else:
            response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)

        file_list += [
            obj["Key"]
            for obj in response.get("Contents", [])
            if not obj["Key"].endswith("/")
        ]

        if response.get("IsTruncated"):
            continuation_token = response.get("NextContinuationToken")
        else:
            break

    contract_list = sorted(file_list)
    files = {}

    for contract in contract_list:
        print(f"Reading {contract}")
        try:
            data = s3_client.get_object(Bucket=bucket, Key=contract)
            contents = data["Body"].read()
            context = contents.decode("utf-8")
            path, filename = os.path.split(contract)
            files[filename] = context
        except:
            pass

    return files


def read_input(local_path=config.LOCAL_PATH):
    if config.READ_MODE == "local":
        files = {}
        for file in os.listdir(local_path):
            full_path = os.path.join(local_path, file)
            file_text = read_local(full_path)
            files[file] = file_text
        return files
    elif config.READ_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 contains_reimbursement(text, page="1"):
    if isinstance(text, dict):
        return page.isdigit() and (
            "%" in text[page]
            or "$" in text[page]
            or "percent " in text[page]
            or "compensation schedule" in text[page].lower()
            or "reimbursement schedule" in text[page].lower()
        )
    elif isinstance(text, str):
        return (
            "%" in text
            or "$" in text
            or "percent " in text
            or "compensation schedule" in text.lower()
            or "reimbursement schedule" in text.lower()
        )
    else:
        print("contains_reimbursement - Invalid data type")


def split_and_write_csv(
    df, n, column_name, output_dir=config.CONSOLIDATED_OUTPUT_DIRECTORY
):
    """
    Splits the DataFrame into N sub-DataFrames based on unique values in a specified column
    and writes each sub-DataFrame to a separate CSV file. Each sub-DataFrame will contain
    roughly the same number of unique values from the specified column.

    Args:
    df (pd.DataFrame): The DataFrame to split.
    n (int): The number of sub-DataFrames to create.
    column_name (str): The column to split the DataFrame by.
    output_dir (str): The directory to save the CSV files to.
    """
    # Create the directory if it does not exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Get unique values and shuffle them to randomize distribution if needed
    unique_values = np.array(df[column_name].unique())
    np.random.shuffle(unique_values)

    # Split the unique values into N parts
    value_groups = np.array_split(unique_values, n)

    # Loop through each group and filter DataFrame by the values in each group
    for i, values in enumerate(value_groups, 1):
        # Filter df for the values in this group
        sub_df = df[df[column_name].isin(values)]
        # Create a filename for each sub-DataFrame
        filename = f"{config.BATCH_ID}-{i}.csv"
        path = os.path.join(output_dir, filename)

        # Write the sub-DataFrame to a CSV file
        sub_df.to_csv(path, index=False)
        print(f"Written {path}")

def contains_lol(text, page='1'):
    if isinstance(text, dict):
        return page.isdigit() and (
            re.search('les+e*o*r of+', text[page].lower())
        )
    elif isinstance(text, str):
        return (
            re.search('les+e*o*r of+', text.lower())
        )
    else:
        print("contains_lol - Invalid data type")
    


def filter_already_processed(input_dict, input_dict_b):
    already_processed_ac, already_processed_b = [], []
    for folder_name in os.listdir(config.OUTPUT_DIRECTORY):
        # AC
        if config.AC_RESULTS_NAME in os.listdir(
            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
        ):
            already_processed_ac.append(folder_name + ".txt")
        # B
        if config.B_RESULTS_NAME in os.listdir(
            os.path.join(config.OUTPUT_DIRECTORY, folder_name)
        ):
            already_processed_b.append(folder_name + ".txt")

    input_dict_ac = {
        key: input_dict[key]
        for key in input_dict.keys()
        if key not in already_processed_ac
    }
    input_dict_b = {
        key: input_dict_b[key]
        for key in input_dict_b.keys()
        if key not in already_processed_b
    }
    return input_dict_ac, input_dict_b


def is_empty(value):
    if pd.isna(value):
        return True
    else:
        empty_values = [None, "", "N/A", "NA", "null", "none", "NaN", np.nan, "nan"]
        return value in empty_values


def find_regex_matches(
    pattern: str, text_dict: dict
) -> tuple[list, list]:  # TODO: Write few unit tests
    """Helper function to find matches and list of pages of a given regex pattern in text_dict.

    Args:
        pattern(str): A particular regex match that with a pattern.
        text_dict (dict): Dictionary keyed by string page-number and valued by actual textract output page

    Returns:
        matches(list): List of match(s)
        page_list(list): List of pages where the match was found

    """
    matches = []
    page_list = []
    for page, text in text_dict.items():
        found = re.findall(pattern, text)
        if found:
            matches.extend(found) # make it append
            page_list.append(page)

    return matches, page_list
    
    # matches = defaultdict(list)
    # for page, text in text_dict.items():
    #     matches_found = re.findall(pattern, text)
    #     if matches_found:
    #         matches[page].extend(matches_found)
    # return list(matches.values()), list(matches.keys())


def chunk_selected_pages(text_dict: dict[str, str], page_list: list) -> str:
    """Combines chunks for pages with regex matches for IRS, NPI, etc.

    Args:
        text_dict (dict[str, str]): A dictionary keyed by string page number and valued
            by string page contents
        page_list (str): List of pages to be chunked for context

    Returns:
        context(str): corresponding chunks of text (input strings delimited by newlines).
    """

    page_list = sorted([int(page_str) for page_str in set(page_list)])
    context = "\n".join([text_dict[str(page_num)] for page_num in page_list])
    return context


def add_hyphen_if_needed(input_str):
    """Checks for specific patterns in the input string and modifies it by adding 
    a '0' at the start and a hyphen at the 3rd index, or just adds a hyphen at 
    the 3rd index if it's a 9-digit number.

    If the input string is exactly 8 characters long, it adds a '0' at the start 
    and inserts a hyphen at the 3rd position. If the string matches the pattern 
    of a 9-digit number (exactly 9 digits), it inserts a hyphen at the 3rd index.
    If the input string matches the pattern 'x-xxxxxxx', it adds '0' at the start.

    All conditions will return a string in the format 'xx-xxxxxxx'.

    Args:
    - input_str (str): The input string (IRS).

    Returns:
    - str: The modified string in the format 'xx-xxxxxxx'.
    
    Example:
    - '1-2345678' → '01-2345678'
    - '12345678' → '01-2345678'
    - '123456789' → '12-3456789'
    """
    # Check if the string matches the pattern 'x-xxxxxxx' (1 character, hyphen, 7 characters)
    if re.fullmatch(r"\d-\d{7}", input_str):
        # Add '0' at the start
        input_str = '0' + input_str
    
    # Check if the string is exactly 'xxxxxxxx' (8 characters)
    elif bool(re.fullmatch(r"\b\d{8}\b", input_str)):
        # Add '0' at the start and insert a hyphen at the 3rd index
        input_str = '0' + input_str
        input_str = input_str[:2] + '-' + input_str[2:]
    
    # If the string matches the pattern of a 9-digit number, add a hyphen at the 3rd index
    elif bool(re.fullmatch(r"\b\d{9}\b", input_str)):
        input_str = input_str[:2] + '-' + input_str[2:]
    
    # Ensure the output is in the format 'xx-xxxxxxx'
    if not bool(re.fullmatch(r"\b\d{2}-\d{7}\b", input_str)):
        input_str = None

    return input_str