doczyai-pipelines/src/postprocessingfuncs.py

import pandas as pd
import re
import difflib

def sanitize_value(value):
    """ Remove brackets from list items and clean the values. """
    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):
    """ Cleans a column by applying an exact match check and updates it to a new column. """
    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):
    """ Applies fuzzy matching to a column in the dataframe and logs changes. """
    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):
    """
    Check and correct misplaced values across specified columns.
    """
    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