hipom_data_mapping/post_process/selection_with_pattern/run.py

# %%
import pandas as pd
import os
import glob
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix
import numpy as np
from utils import BertEmbedder, cosine_similarity_chunked
from fuzzywuzzy import fuzz

##################
# global parameters
DIAGNOSTIC = False
THRESHOLD = 0.85
FUZZY_SIM_THRESHOLD=95
checkpoint_directory = "../../train/classification_bert_desc"

###################
# %%
# helper functions
class Embedder():
    input_df: pd.DataFrame
    fold: int

    def __init__(self, input_df):
        self.input_df = input_df


    def make_embedding(self, checkpoint_path):

        def generate_input_list(df):
            input_list = []
            for _, row in df.iterrows():
                desc = f"{row['tag_description']}"
                element = f"{desc}"
                input_list.append(element)
            return input_list

        # prepare reference embed
        train_data = list(generate_input_list(self.input_df))
        # Define the directory and the pattern
        embedder = BertEmbedder(train_data, checkpoint_path)
        embedder.make_embedding(batch_size=64)
        return embedder.embeddings.to('cpu')



# the purpose of this function is to group the strings that are similar to each other
# we need to form related groups of inputs
def group_similar_strings(obj_list, threshold=80):
    groups = []
    processed_strings = set()  # To keep track of already grouped strings

    for obj in obj_list:
        # tuple is (idx, string)
        if obj in processed_strings:
            continue
        # Find all strings similar to the current string above the threshold
        similar_strings = [s for s in obj_list if s[1] != obj[1] and fuzz.ratio(obj[1], s[1]) >= threshold]
        # Add the original string to the similar group
        similar_group = [obj] + similar_strings
        # Mark all similar strings as processed
        processed_strings.update(similar_group)
        # Add the group to the list of groups
        groups.append(similar_group)
    return groups

# this function takes in groups of related terms and create candidate entries
def make_candidates(groups):
    candidates = []
    for group in groups:
        first_tuple = group[0]
        # string_of_tuple = first_tuple[1]
        id_of_tuple = first_tuple[0]
        candidates.append(id_of_tuple)
    return candidates


# the selection function takes in the full cos_sim_matrix then subsets the
# matrix according to the test_candidates_mask and train_candidates_mask that we
# give it
# it returns the most likely source candidate index and score among the source
# candidate list - aka it returns a local idx
def selection(cos_sim_matrix, source_mask, target_mask):
    # subset_matrix = cos_sim_matrix[condition_source]
    # except we are subsetting 2D matrix (row, column)
    subset_matrix = cos_sim_matrix[np.ix_(source_mask, target_mask)]
    # we select top-k here
    # Get the indices of the top-k maximum values along axis 1
    top_k = 1
    top_k_indices = np.argsort(subset_matrix, axis=1)[:, -top_k:]  # Get indices of top k values
    
    # Get the values of the top 5 maximum scores
    top_k_values = np.take_along_axis(subset_matrix, top_k_indices, axis=1)
    
    # Calculate the average of the top-k scores along axis 1
    y_scores = np.mean(top_k_values, axis=1)
    max_idx = np.argmax(y_scores)
    max_score = y_scores[max_idx]
    

    return max_idx, max_score



####################
# global level
# %%
# obtain the full mdm_list
data_path = '../../data_import/exports/data_mapping_mdm.csv'
full_df = pd.read_csv(data_path, skipinitialspace=True)
full_mdm_pattern_list = sorted(list((set(full_df['pattern']))))


#####################
# fold level

def run_selection(fold):

    # set the fold
    # import test data
    data_path = f"../../train/mapping_pattern/mapping_prediction/exports/result_group_{fold}.csv"
    df = pd.read_csv(data_path, skipinitialspace=True)
    df['p_pattern'] = df['p_thing'] + " " + df['p_property']

    # get target data
    data_path = f"../../data_preprocess/exports/dataset/group_{fold}/train_all.csv"
    train_df = pd.read_csv(data_path, skipinitialspace=True)

    # generate your embeddings
    # checkpoint_directory defined at global level
    directory = os.path.join(checkpoint_directory, f'checkpoint_fold_{fold}')
    # Use glob to find matching paths
    # path is usually checkpoint_fold_1/checkpoint-<step number>
    # we are guaranteed to save only 1 checkpoint from training
    pattern = 'checkpoint-*'
    checkpoint_path = glob.glob(os.path.join(directory, pattern))[0]

    # we can generate the train embeddings once and re-use for every ship
    train_embedder = Embedder(input_df=train_df)
    train_embeds = train_embedder.make_embedding(checkpoint_path)


    # create global_answer array
    # the purpose of this array is to track the classification state at the global
    # level
    global_answer = np.zeros(len(df), dtype=bool)


    #############################
    # ship level
    # we have to split into per-ship analysis
    ships_list = sorted(list(set(df['ships_idx'])))
    for ship_idx in ships_list:
        # ship_idx = 1001  # choose an example ship

        ship_df = df[df['ships_idx'] == ship_idx]
        # required to map local ship_answer array to global_answer array
        map_local_index_to_global_index = ship_df.index.to_numpy()
        ship_df = df[df['ships_idx'] == ship_idx].reset_index(drop=True)

        # generate new embeddings for each ship
        test_embedder = Embedder(input_df=ship_df)
        test_embeds = test_embedder.make_embedding(checkpoint_path)

        # generate the cosine sim matrix
        cos_sim_matrix = cosine_similarity_chunked(test_embeds, train_embeds, chunk_size=8).cpu().numpy()

        ##############################
        # selection level
        # The general idea:
        # step 1: keep only pattern generations that belong to mdm list
        # -> this removes totally wrong datasets that mapped to totally wrong things
        # step 2: loop through the mdm list and isolate data in both train and test that
        # belong to the same pattern class
        # -> this is more tricky, because we have non-mdm mapping to correct classes
        # -> so we have to find which candidate is most similar to the training data

        # it is very tricky to keep track of classification across multiple stages so we
        # will use a boolean answer list

        # initialize the local answer list
        ship_answer_list = np.ones(len(ship_df), dtype=bool)

        ###########
        # STEP 1
        # we want to loop through the generated class labels and find which ones match
        # our pattern list

        pattern_match_mask = ship_df['p_pattern'].apply(lambda x: x in full_mdm_pattern_list).to_numpy()
        # we assign only those that are False to our answer list
        # right now the 2 arrays are basically equal
        ship_answer_list[~pattern_match_mask] = False

        ###########
        # STEP 2
        # we now go through each class found in our generated set

        # we want to identify per-ship mdm classes
        ship_predicted_classes = sorted(set(ship_df['p_pattern'][pattern_match_mask].to_list()))

        # this function performs the selection given a class
        # it takes in the cos_sim_matrix
        # it returns the selection by mutating the answer_list
        # it sets all relevant idxs to False initially, then sets the selected values to True
        def selection_for_class(select_class, cos_sim_matrix, answer_list):

            # separate the global variable from function variable
            answer_list = answer_list.copy()
            sample_df = ship_df[ship_df['p_pattern'] == select_class]

            # we need to set all idx of chosen entries as False in answer_list
            selected_idx_list = sample_df.index.to_list()
            answer_list[selected_idx_list] = False

            # basic assumption check

            # group related inputs by description similarity
            string_list = sample_df['tag_description'].to_list()
            index_list = sample_df.index.to_list()
            obj_list = list(zip(index_list, string_list))
            # groups is a list of list, where each list is composed of a
            # (idx, string) tuple
            groups = group_similar_strings(obj_list, threshold=FUZZY_SIM_THRESHOLD)

            # generate the masking arrays for both test and train embeddings
            # we select a tuple from each group, and use that as a candidate for selection
            test_candidates = make_candidates(groups)
            test_candidates_mask = np.zeros(len(ship_df), dtype=bool)
            test_candidates_mask[test_candidates] = True
            # we make candidates to compare against in the data sharing the same class
            train_candidates_mask = (train_df['pattern'] == select_class).to_numpy()

            # perform selection
            # it returns the group index that is most likely
            max_idx, max_score = selection(cos_sim_matrix, test_candidates_mask, train_candidates_mask)

            # consolidate all idx's in the same group
            chosen_group = groups[max_idx]
            chosen_idx_list = [tuple[0] for tuple in chosen_group]


            # before doing this, we have to use the max_score and evaluate if its close enough
            if max_score > THRESHOLD:
                answer_list[chosen_idx_list] = True

            return answer_list 


        # we choose one mdm class
        for select_class in ship_predicted_classes:
            ship_answer_list = selection_for_class(select_class, cos_sim_matrix, ship_answer_list)

        # we want to write back to global_answer
        # first we convert local indices to global indices
        local_indices = np.where(ship_answer_list)[0]
        global_indices = map_local_index_to_global_index[local_indices]
        global_answer[global_indices] = True


        if DIAGNOSTIC:
            # evaluation at per-ship level
            y_true = ship_df['MDM'].to_list()
            y_pred = ship_answer_list

            # Compute metrics
            accuracy = accuracy_score(y_true, y_pred)
            f1 = f1_score(y_true, y_pred, average='macro')
            precision = precision_score(y_true, y_pred, average='macro')
            recall = recall_score(y_true, y_pred, average='macro')

            # Print the results
            print(f'Accuracy: {accuracy:.5f}')
            print(f'F1 Score: {f1:.5f}')
            print(f'Precision: {precision:.5f}')
            print(f'Recall: {recall:.5f}')



    y_true = df['MDM'].to_list()
    y_pred = global_answer

    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
    print(f"tp: {tp}")
    print(f"tn: {tn}")
    print(f"fp: {fp}")
    print(f"fn: {fn}")

    # Compute metrics
    accuracy = accuracy_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred, average='macro')
    precision = precision_score(y_true, y_pred, average='macro')
    recall = recall_score(y_true, y_pred, average='macro')

    # Print the results
    print(f'Accuracy: {accuracy:.5f}')
    print(f'F1 Score: {f1:.5f}')
    print(f'Precision: {precision:.5f}')
    print(f'Recall: {recall:.5f}')


# %%
for fold in [1,2,3,4,5]:
    print(f'Perform selection for fold {fold}')
    run_selection(fold)


# %%