hipom_data_mapping/train/class_number_tokens/classification_prediction/predict.py

# %%

# from datasets import load_from_disk
import os
import glob

os.environ['NCCL_P2P_DISABLE'] = '1'
os.environ['NCCL_IB_DISABLE'] = '1'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"

import torch
from torch.utils.data import DataLoader

from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
)
import evaluate
import re
import numpy as np
import pandas as pd
# import matplotlib.pyplot as plt
from datasets import Dataset, DatasetDict

from tqdm import tqdm

torch.set_float32_matmul_precision('high')


BATCH_SIZE = 128

# %%

# we need to create the mdm_list
# import the full mdm-only file
data_path = '../../../data_import/exports/data_mapping_mdm.csv'
full_df = pd.read_csv(data_path, skipinitialspace=True)
# rather than use pattern, we use the real thing and property
# mdm_list = sorted(list((set(full_df['pattern']))))
thing_property = full_df['thing'] + full_df['property']
thing_property = thing_property.to_list()
mdm_list = sorted(list(set(thing_property)))


# %%
id2label = {}
label2id = {}
for idx, val in enumerate(mdm_list):
    id2label[idx] = val
    label2id[val] = idx

# %%
def substitute_and_append_digits(s):
    """
    Finds all digit groups in a string, substitutes them with a <digit> placeholder, 
    and appends the extracted digit groups at the end of the string flanked by <digit> markers.

    Args:
        s (str): The input string.

    Returns:
        str: The transformed string.
    """
    # Find all digit groups in the string
    digit_groups = re.findall(r'\d+', s)
    
    # Substitute digit groups with <digit> placeholder
    substituted_string = re.sub(r'\d+', '<DIGIT>', s)
    
    # Append extracted digit groups to the end of the string
    appended_digits = ''.join([f'<DIGIT>{group}<DIGIT>' for group in digit_groups])
    
    return substituted_string + appended_digits



# outputs a list of dictionaries
# processes dataframe into lists of dictionaries
# each element maps input to output
# input: tag_description
# output: class label
def process_df_to_dict(df, mdm_list):
    output_list = []
    for _, row in df.iterrows():
        processed_desc = substitute_and_append_digits(row['tag_description'])
        desc = f"<DESC>{processed_desc}<DESC>"
        unit = f"<UNIT>{row['unit']}<UNIT>"

        pattern = f"{row['thing'] + row['property']}"
        try:
            index = mdm_list.index(pattern)
        except ValueError:
            index = -1
        element = {
            'text' : f"{desc}{unit}",
            'label': index,
        }
        output_list.append(element)

    return output_list


def create_dataset(fold, mdm_list):
    data_path = f"../../../data_preprocess/exports/dataset/group_{fold}/test_all.csv"
    test_df = pd.read_csv(data_path, skipinitialspace=True)
    # uncomment for mdm
    # we only use the mdm subset
    test_df = test_df[test_df['MDM']].reset_index(drop=True)

    test_dataset = Dataset.from_list(process_df_to_dict(test_df, mdm_list))

    return test_dataset


# %%

# function to perform training for a given fold
def test(fold):

    test_dataset = create_dataset(fold, mdm_list)

    # prepare tokenizer

    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-*'
    model_checkpoint = glob.glob(os.path.join(checkpoint_directory, pattern))[0]

    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, return_tensors="pt", clean_up_tokenization_spaces=True)
    # Define additional special tokens
    additional_special_tokens = ["<THING_START>", "<THING_END>", "<PROPERTY_START>", "<PROPERTY_END>", "<DESC>", "<UNIT>", "<DIGIT>"]
    # Add the additional special tokens to the tokenizer
    tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens})

    # %%
    # compute max token length
    max_length = 0
    for sample in test_dataset['text']:
        # Tokenize the sample and get the length
        input_ids = tokenizer(sample, truncation=False, add_special_tokens=True)["input_ids"]
        length = len(input_ids)
        
        # Update max_length if this sample is longer
        if length > max_length:
            max_length = length

    print(max_length)

    # %%

    max_length = 128

    # given a dataset entry, run it through the tokenizer
    def preprocess_function(example):
        input = example['text']
        # text_target sets the corresponding label to inputs
        # there is no need to create a separate 'labels'
        model_inputs = tokenizer(
            input,
            max_length=max_length,
            # truncation=True,
            padding='max_length'
        )
        return model_inputs

    # map maps function to each "row" in the dataset
    # aka the data in the immediate nesting
    datasets = test_dataset.map(
        preprocess_function,
        batched=True,
        num_proc=8,
        remove_columns="text",
    )


    datasets.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])

    # %% temp
    # tokenized_datasets['train'].rename_columns()

    # %%
    # create data collator

    # data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="max_length")

    # %%
    # compute metrics
    # metric = evaluate.load("accuracy")
    # 
    # 
    # def compute_metrics(eval_preds):
    #     preds, labels = eval_preds
    #     preds = np.argmax(preds, axis=1)
    #     return metric.compute(predictions=preds, references=labels)

    model = AutoModelForSequenceClassification.from_pretrained(
        model_checkpoint,
        num_labels=len(mdm_list),
        id2label=id2label,
        label2id=label2id)
    # important! after extending tokens vocab
    model.resize_token_embeddings(len(tokenizer))

    model = model.eval()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    pred_labels = []
    actual_labels = []


    dataloader = DataLoader(datasets, batch_size=BATCH_SIZE, shuffle=False)
    for batch in tqdm(dataloader):
            # Inference in batches
            input_ids = batch['input_ids']
            attention_mask = batch['attention_mask']
            # save labels too
            actual_labels.extend(batch['label'])
            

            # Move to GPU if available
            input_ids = input_ids.to(device)
            attention_mask = attention_mask.to(device)

            # Perform inference
            with torch.no_grad():
                logits = model(
                    input_ids,
                    attention_mask).logits
                predicted_class_ids = logits.argmax(dim=1).to("cpu")
                pred_labels.extend(predicted_class_ids)

    pred_labels = [tensor.item() for tensor in pred_labels]


    # %%
    from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix
    y_true = actual_labels
    y_pred = pred_labels

    # Compute metrics
    accuracy = accuracy_score(y_true, y_pred)
    average_parameter = 'weighted'
    zero_division_parameter = 0
    f1 = f1_score(y_true, y_pred, average=average_parameter, zero_division=zero_division_parameter)
    precision = precision_score(y_true, y_pred, average=average_parameter, zero_division=zero_division_parameter)
    recall = recall_score(y_true, y_pred, average=average_parameter, zero_division=zero_division_parameter)



    with open("output.txt", "a") as f:

        print('*' * 80, file=f)
        print(f'Fold: {fold}', file=f)
        # Print the results
        print(f'Accuracy: {accuracy:.5f}', file=f)
        print(f'F1 Score: {f1:.5f}', file=f)
        print(f'Precision: {precision:.5f}', file=f)
        print(f'Recall: {recall:.5f}', file=f)

    # export result
    data_path = f"../../../data_preprocess/exports/dataset/group_{fold}/test_all.csv"
    df = pd.read_csv(data_path, skipinitialspace=True)
    # uncomment if you want to predict for all
    # df = df[df['MDM']].reset_index(drop=True)

    label_list = [id2label[id] for id in pred_labels]
    df_out = pd.DataFrame({
        'class_prediction': pd.Series(label_list) 
    })
    df = pd.concat([df, df_out], axis=1)

    # we can save the t5 generation output here
    df.to_csv(f"exports/result_group_{fold}.csv", index=False)



# %%
# reset file before writing to it
with open("output.txt", "w") as f:
    print('', file=f)

for fold in [1,2,3,4,5]:
    test(fold)