# %% # 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 placeholder, and appends the extracted digit groups at the end of the string flanked by 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 placeholder substituted_string = re.sub(r'\d+', '', s) # Append extracted digit groups to the end of the string appended_digits = ''.join([f'{group}' 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"{processed_desc}" unit = f"{row['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- # 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 = ["", "", "", "", "", "", ""] # 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)