hipom_data_mapping/post_process/similarity_classifier/run.py

135 lines
4.4 KiB
Python
Raw Normal View History

# %%
import pandas as pd
from utils import Retriever, cosine_similarity_chunked
import os
import glob
import numpy as np
from tqdm import tqdm
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix
##################################################
# helper functions
# the following function takes in a full cos_sim_matrix
# condition_source: boolean selectors of the source embedding
# condition_target: boolean selectors of the target embedding
def find_closest(cos_sim_matrix, condition_source, condition_target):
# subset_matrix = cos_sim_matrix[condition_source]
# except we are subsetting 2D matrix (row, column)
subset_matrix = cos_sim_matrix[np.ix_(condition_source, condition_target)]
# we select top k here
# Get the indices of the top k maximum values along axis 1
top_k = 3
top_k_indices = np.argsort(subset_matrix, axis=1)[:, -top_k:] # Get indices of top k values
# note that top_k_indices is a nested list because of the 2d nature of the matrix
# the result is flipped
top_k_indices[0] = top_k_indices[0][::-1]
# Get the values of the top 5 maximum scores
top_k_values = np.take_along_axis(subset_matrix, top_k_indices, axis=1)
return top_k_indices, top_k_values
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"<DESC>{row['tag_description']}<DESC>"
unit = f"<UNIT>{row['unit']}<UNIT>"
element = f"{desc}{unit}"
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
retriever_train = Retriever(train_data, checkpoint_path)
retriever_train.make_embedding(batch_size=64)
return retriever_train.embeddings.to('cpu')
def run_similarity_classifier(fold):
data_path = f'../../train/mapping_pattern/mapping_prediction/exports/result_group_{fold}.csv'
test_df = pd.read_csv(data_path, skipinitialspace=True)
data_path = f"../../data_preprocess/exports/dataset/group_{fold}/train_all.csv"
train_df = pd.read_csv(data_path, skipinitialspace=True)
checkpoint_directory = "../../train/classification_bert_complete_desc_unit"
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]
train_embedder = Embedder(input_df=train_df)
train_embeds = train_embedder.make_embedding(checkpoint_path)
test_embedder = Embedder(input_df=test_df)
test_embeds = test_embedder.make_embedding(checkpoint_path)
def compute_top_k(select_idx):
condition_source = test_df['tag_description'] == test_df[test_df.index == select_idx]['tag_description'].tolist()[0]
condition_target = np.ones(train_embeds.shape[0], dtype=bool)
_, top_k_values = find_closest(
cos_sim_matrix=cos_sim_matrix,
condition_source=condition_source,
condition_target=condition_target)
return top_k_values[0][0]
# test embeds are inputs since we are looking back at train data
cos_sim_matrix = cosine_similarity_chunked(test_embeds, train_embeds, chunk_size=1024).cpu().numpy()
sim_list = []
for select_idx in tqdm(test_df.index):
top_sim_value = compute_top_k(select_idx)
sim_list.append(top_sim_value)
# analysis 1: using threshold to perform find-back prediction success
threshold = 0.90
predict_list = [ elem > threshold for elem in sim_list ]
y_true = test_df['MDM'].to_list()
y_pred = predict_list
# Compute metrics
accuracy = accuracy_score(y_true, y_pred)
f1 = f1_score(y_true, y_pred)
precision = precision_score(y_true, y_pred)
recall = recall_score(y_true, y_pred)
# 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]:
run_similarity_classifier(fold)