257 lines
7.3 KiB
Python
257 lines
7.3 KiB
Python
# %%
|
|
|
|
# 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 re
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from transformers import (
|
|
AutoTokenizer,
|
|
AutoModelForSequenceClassification,
|
|
AutoModel,
|
|
DataCollatorWithPadding,
|
|
)
|
|
import evaluate
|
|
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 = 256
|
|
|
|
# %%
|
|
# construct the target id list
|
|
# data_path = '../../../esAppMod_data_import/train.csv'
|
|
data_path = '../esAppMod_data_import/train.csv'
|
|
train_df = pd.read_csv(data_path, skipinitialspace=True)
|
|
# rather than use pattern, we use the real thing and property
|
|
entity_ids = train_df['entity_id'].to_list()
|
|
target_id_list = sorted(list(set(entity_ids)))
|
|
|
|
|
|
# %%
|
|
id2label = {}
|
|
label2id = {}
|
|
for idx, val in enumerate(target_id_list):
|
|
id2label[idx] = val
|
|
label2id[val] = idx
|
|
|
|
|
|
# introduce pre-processing functions
|
|
def preprocess_text(text):
|
|
# 1. Make all uppercase
|
|
text = text.lower()
|
|
|
|
# Substitute digits with '#'
|
|
# text = re.sub(r'\d+', '#', text)
|
|
|
|
# standardize spacing
|
|
text = re.sub(r'\s+', ' ', text).strip()
|
|
|
|
return text
|
|
|
|
|
|
|
|
|
|
# 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):
|
|
output_list = []
|
|
for _, row in df.iterrows():
|
|
desc = row['mention']
|
|
desc = preprocess_text(desc)
|
|
index = row['entity_id']
|
|
element = {
|
|
'text' : desc,
|
|
'label': label2id[index], # ensure labels starts from 0
|
|
}
|
|
output_list.append(element)
|
|
|
|
return output_list
|
|
|
|
|
|
def create_dataset():
|
|
# train
|
|
data_path = '../esAppMod_data_import/test.csv'
|
|
test_df = pd.read_csv(data_path, skipinitialspace=True)
|
|
|
|
|
|
# combined_data = DatasetDict({
|
|
# 'train': Dataset.from_list(process_df_to_dict(train_df)),
|
|
# })
|
|
return Dataset.from_list(process_df_to_dict(test_df))
|
|
|
|
|
|
|
|
# %%
|
|
|
|
def test():
|
|
|
|
test_dataset = create_dataset()
|
|
|
|
# prepare tokenizer
|
|
|
|
MODEL_NAME = 'prajjwal1/bert-small' # 'prajjwal1/bert-small' 'bert-base-cased' 'distilbert-base-cased'
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, return_tensors="pt", clean_up_tokenization_spaces=True)
|
|
# Define additional special tokens
|
|
# additional_special_tokens = ["<THING_START>", "<THING_END>", "<PROPERTY_START>", "<PROPERTY_END>", "<NAME>", "<DESC>", "<SIG>", "<UNIT>", "<DATA_TYPE>"]
|
|
# 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,
|
|
# 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'])
|
|
|
|
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
|
bert_model = AutoModel.from_pretrained(MODEL_NAME)
|
|
|
|
class BertForClassificationAndTriplet(nn.Module):
|
|
def __init__(self, bert_model, num_classes):
|
|
super().__init__()
|
|
self.bert = bert_model
|
|
self.classifier = nn.Linear(bert_model.config.hidden_size, num_classes)
|
|
|
|
def forward(self, input_ids, attention_mask=None):
|
|
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
|
|
cls_embeddings = outputs.last_hidden_state[:, 0, :] # CLS token
|
|
logits = self.classifier(cls_embeddings)
|
|
return cls_embeddings, logits
|
|
|
|
model = BertForClassificationAndTriplet(bert_model, num_classes=len(label2id))
|
|
state_dict = torch.load('./checkpoint/classification.pt')
|
|
model.load_state_dict(state_dict)
|
|
|
|
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, collate_fn=data_collator)
|
|
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['labels'])
|
|
|
|
|
|
# Move to GPU if available
|
|
input_ids = input_ids.to(device)
|
|
attention_mask = attention_mask.to(device)
|
|
|
|
# Perform inference
|
|
with torch.no_grad():
|
|
cls, logits = model(
|
|
input_ids,
|
|
attention_mask)
|
|
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("results/output.txt", "a") as f:
|
|
|
|
print('*' * 80, 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
|
|
label_list = [id2label[id] for id in pred_labels]
|
|
df = pd.DataFrame({
|
|
'class_prediction': pd.Series(label_list)
|
|
})
|
|
|
|
# we can save the t5 generation output here
|
|
df.to_csv(f"results/classify.csv", index=False)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# %%
|
|
# reset file before writing to it
|
|
with open("results/output.txt", "w") as f:
|
|
print('', file=f)
|
|
test()
|