doczyai-pipelines/src/postprocessingfuncs.py

import pandas as pd
import re
import difflib

def sanitize_value(value):
    if pd.isna(value):
        return value
    if isinstance(value, list):
        value = ', '.join(str(v) for v in value)
    if isinstance(value, str):
        value = value.strip('[]')
        value = ', '.join([item.strip(" '") for item in value.split(',')])
    return value

def exact_match(val, valid_values):
    val = val.strip().upper()
    for valid_val in valid_values:
        if val == valid_val.upper():
            return valid_val
    return None

def clean_columns_combined(df, column_name, valid_values, new_column_name):
    changes = {}
    df[column_name] = df[column_name].apply(sanitize_value)
    original_values = df[column_name].unique()

    def update_column(entry):
        if pd.notna(entry):
            terms = entry.split(',')
            for term in terms:
                match = exact_match(term, valid_values)
                if match:
                    return match
        return None
    
    df[new_column_name] = df[column_name].apply(update_column)
    cleaned_values = df[new_column_name].unique()
    return original_values, cleaned_values, changes

def get_closest_match(val, valid_values, similarity_threshold=0.7):
    if pd.isna(val):
        return None
    val = val.strip().upper()
    matches = difflib.get_close_matches(val, [v.upper() for v in valid_values], n=1, cutoff=similarity_threshold)
    return matches[0] if matches else None

def clean_columns_combined_fuzzy(df, column_name, valid_values, threshold):
    changes = {}
    df[column_name] = df[column_name].apply(sanitize_value)
    original_values = df[column_name].unique()
    
    def log_and_clean(entry):
        if pd.notna(entry):
            words = entry.split(',')
            cleaned_words = []
            for word in words:
                cleaned_word = get_closest_match(word.strip(), valid_values, similarity_threshold=threshold)
                if cleaned_word and word.strip().upper() != cleaned_word:
                    changes[word.strip()] = cleaned_word
                cleaned_words.append(cleaned_word if cleaned_word else word.strip())
            return ', '.join(cleaned_words)
        return None

    df[column_name] = df[column_name].apply(log_and_clean)
    cleaned_values = df[column_name].unique()
    return original_values, cleaned_values, changes

def extract_page_number(page_text):
    match = re.search(r'Pages\s+(\d+)-\d+', page_text)
    if match:
        return match.group(1)
    else:
        return page_text

def clean_pagenumbers(df):
    df['page_num'] = df['page_num'].apply(extract_page_number)
    return df

def correct_misplaced_values(df, columns, valid_values_dict):
    for index, row in df.iterrows():
        for col in columns:
            if pd.notna(row[col]):
                terms = row[col].split(',')
                for term in terms:
                    term = term.strip()
                    for target_col, valid_values in valid_values_dict.items():
                        if target_col != col:
                            match = exact_match(term, valid_values)
                            if match:
                                if pd.isna(row[target_col]) or not row[target_col].strip():
                                    df.at[index, target_col] = match
                                    df.at[index, col] = None
                                else:
                                    current_value = row[target_col].strip()
                                    if get_closest_match(match, [current_value], 0.8) is None:
                                        df.at[index, 'Corrected_' + target_col] = f"Found {term} in {col} cell"
                                        df.at[index, col] = None
    return df

def filter_service_column(df):
    keywords = [
        'LIABILITY', 'RISK', 'LOBBYING', 'DAMAGES', 'CONFIDENTIALITY', 'AUDIT', 'INTEREST'
    ]
    pattern = '|'.join(keywords)
    df = df[~df['SERVICE'].str.contains(pattern, case=False, na=False)]
    return df

def move_percentage_to_rate(df):
    def move_percentage(value):
        if isinstance(value, str) and '%' in value:
            return True
        return False
    
    df['REIMBURSEMENT_RATE'] = df.apply(lambda row: row['REIMBURSEMENT_FLAT_FEE'] if move_percentage(row['REIMBURSEMENT_FLAT_FEE']) else row['REIMBURSEMENT_RATE'], axis=1)
    df['REIMBURSEMENT_FLAT_FEE'] = df.apply(lambda row: None if move_percentage(row['REIMBURSEMENT_FLAT_FEE']) else row['REIMBURSEMENT_FLAT_FEE'], axis=1)
    return df

def move_large_numbers_to_flat_fee(df):
    def move_large_number(value):
        if isinstance(value, str) and value.isdigit() and int(value) > 100:
            return True
        return False
    
    df['REIMBURSEMENT_FLAT_FEE'] = df.apply(lambda row: row['REIMBURSEMENT_RATE'] if move_large_number(row['REIMBURSEMENT_RATE']) else row['REIMBURSEMENT_FLAT_FEE'], axis=1)
    df['REIMBURSEMENT_RATE'] = df.apply(lambda row: None if move_large_number(row['REIMBURSEMENT_RATE']) else row['REIMBURSEMENT_RATE'], axis=1)
    return df

def set_rate_to_zero_if_not_covered(df):
    def check_and_set_rate(row):
        if pd.notna(row['FULL_METHODOLOGY']) and re.search(r'not covered', row['FULL_METHODOLOGY'], re.IGNORECASE):
            return 0
        return row['REIMBURSEMENT_RATE']
    
    df['REIMBURSEMENT_RATE'] = df.apply(check_and_set_rate, axis=1)
    return df

def adjust_reimbursement_rate(df):
    def adjust_rate(row):
        if pd.notna(row['FULL_METHODOLOGY']) and re.search(r'case insensitive', row['FULL_METHODOLOGY'], re.IGNORECASE):
            if pd.notna(row['REIMBURSEMENT_RATE']) and isinstance(row['REIMBURSEMENT_RATE'], (int, float)):
                return 100 - row['REIMBURSEMENT_RATE']
        return row['REIMBURSEMENT_RATE']
    
    df['REIMBURSEMENT_RATE'] = df.apply(adjust_rate, axis=1)
    return df