import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler
import pickle
import json
import matplotlib.pyplot as plt
from glob import glob
import time
import copy
from tqdm import tqdm
import re
from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig, BertTokenizer
from data import ZuCo_dataset
from model_decoding import BrainTranslator, BrainTranslatorNaive
from metrics import compute_metrics
from config import get_config
def eval_model(dataloaders, device, tokenizer, criterion, model, output_all_results_path='./results/temp.txt'):
# modified from: https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
model.eval() # Set model to evaluate mode
running_loss = 0.0
# Iterate over data.
# sample_count = 0
target_tokens_list = []
target_string_list = []
pred_tokens_list = []
pred_string_list = []
with open(output_all_results_path, 'w') as f:
# count=0
for input_embeddings, seq_len, input_masks, input_mask_invert, target_ids, target_mask, sentiment_labels, sent_level_EEG in \
dataloaders['test']:
# count+=1
# if count>5:
# break
# load in batch
input_embeddings_batch = input_embeddings.to(device).float()
input_masks_batch = input_masks.to(device)
target_ids_batch = target_ids.to(device)
input_mask_invert_batch = input_mask_invert.to(device)
if intput_noise:
input_embeddings_batch=torch.rand_like(input_embeddings_batch)
# target_tokens = tokenizer.convert_ids_to_tokens(target_ids_batch, skip_special_tokens = True)
target_string = tokenizer.batch_decode(target_ids_batch, skip_special_tokens=True)
# print('target ids tensor:',target_ids_batch[0])
# print('target ids:',target_ids_batch[0].tolist())
# print('+' * 100)
# print('target tokens:',target_tokens)
# print('target string:', target_string)
# add to list for later calculate bleu metric
# target_tokens_list.append([target_tokens])
target_string_list.extend(target_string)
"""replace padding ids in target_ids with -100"""
target_ids_batch[target_ids_batch == tokenizer.pad_token_id] = -100
if not teacher_forcing:
predictions = model.generate(input_embeddings_batch, input_masks_batch, input_mask_invert_batch,
target_ids_batch,
max_length=100,
num_beams=5, do_sample=False, repetition_penalty=5.0,
# num_beams=5,encoder_no_repeat_ngram_size =1,
# do_sample=True, top_k=15,temperature=0.5,num_return_sequences=5,
# early_stopping=True
)
# predicted_string=predicted_string.squeeze()
# print(f'predictions:{predictions}')
# print(f'predicted_string:{predicted_string}')
#
# print(f'predicted_string:{predicted_string}')
else:
seq2seqLMoutput = model(input_embeddings_batch, input_masks_batch, input_mask_invert_batch,
target_ids_batch)
logits = seq2seqLMoutput.logits # bs*seq_len*voc_sz
probs = logits.softmax(dim=-1)
values, predictions = probs.topk(1)
predictions = torch.squeeze(predictions, dim=-1)
# print(f'predictions:{predictions} predictions shape:{predictions.shape}')
predicted_string = tokenizer.batch_decode(predictions, skip_special_tokens=True, )
# print(f'predicted_string:{predicted_string}')
# start = predicted_string.find("[CLS]") + len("[CLS]")
# end = predicted_string.find("[SEP]")
# predicted_string = predicted_string[start:end]
# predicted_string=merge_consecutive_duplicates(predicted_string,'。')
# predictions=tokenizer.encode(predicted_string)
for str_id in range(len(target_string)):
f.write(f'start################################################\n')
f.write(f'Predicted: {predicted_string[str_id]}\n')
f.write(f'True: {target_string[str_id]}\n')
f.write(f'end################################################\n\n\n')
# convert to int list
# predictions = predictions.tolist()
# truncated_prediction = []
# for t in predictions:
# if t != tokenizer.eos_token_id:
# truncated_prediction.append(t)
# else:
# break
# pred_tokens = tokenizer.convert_ids_to_tokens(truncated_prediction, skip_special_tokens = True)
# pred_tokens_list.append(pred_tokens)
pred_string_list.extend(predicted_string)
# sample_count += 1
# print('predicted tokens:',pred_tokens)
# print('predicted string:',predicted_string)
# print('-' * 100)
# print(f'pred_string_list:{pred_string_list}')
# print(f'target_string_list:{target_string_list}')
metrics_results=compute_metrics(pred_string_list,target_string_list)
print(f'teacher_forcing{teacher_forcing} intput_noise{intput_noise}')
print(metrics_results)
print(output_all_results_path)
print(output_all_metrics_results_path)
with open(output_all_metrics_results_path, "w") as json_file:
json.dump(metrics_results, json_file, indent=4, ensure_ascii=False)
if __name__ == '__main__':
home_directory = os.path.expanduser("~")
''' get args'''
args = get_config('eval_decoding')
''' load training config'''
training_config = json.load(open(args['config_path']))
batch_size = 1
subject_choice = training_config['subjects']
print(f'[INFO]subjects: {subject_choice}')
eeg_type_choice = training_config['eeg_type']
print(f'[INFO]eeg type: {eeg_type_choice}')
bands_choice = training_config['eeg_bands']
print(f'[INFO]using bands: {bands_choice}')
dataset_setting = 'unique_sent'
task_name = training_config['task_name']
model_name = training_config['model_name']
# model_name = 'BrainTranslator'
# model_name = 'BrainTranslatorNaive'
# teacher_forcing = True
# {'wer': 0.7980769276618958, 'rouge1_fmeasure': 23.912235260009766, 'rouge1_precision': 24.66936492919922, 'rouge1_recall': 23.318071365356445, 'rouge2_fmeasure': 6.851282119750977, 'rouge2_precision': 6.962162017822266, 'rouge2_recall': 6.751219272613525, 'rougeL_fmeasure': 22.912235260009766, 'rougeL_precision': 23.61673355102539, 'rougeL_recall': 22.36568832397461, 'rougeLsum_fmeasure': 22.912235260009766, 'rougeLsum_precision': 23.61673355102539, 'rougeLsum_recall': 22.36568832397461, 'bleu-1': 0.23883000016212463, 'bleu-2': 0.13888777792453766, 'bleu-3': 0.0, 'bleu-4': 0.0}
teacher_forcing = eval(args['tf'])
intput_noise = eval(args['noise'])
print(f'teacher_forcing{teacher_forcing} intput_noise{intput_noise}')
output_all_results_path = (f'./results/{task_name}-{model_name}{"-teacher_forcing" if teacher_forcing else ""}{"-intput_noise" if intput_noise else ""}-all_decoding_results.txt')
output_all_metrics_results_path = output_all_results_path.replace('txt', 'json')
''' set random seeds '''
seed_val = 312
np.random.seed(seed_val)
torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val)
''' set up device '''
# use cuda
if torch.cuda.is_available():
dev = args['cuda']
else:
dev = "cpu"
# CUDA_VISIBLE_DEVICES=0,1,2,3
device = torch.device(dev)
print(f'[INFO]using device {dev}')
''' set up dataloader '''
whole_dataset_dicts = []
if 'task1' in task_name:
dataset_path_task1 = 'datasets/ZuCo/task1-SR/pickle/task1-SR-dataset.pickle'
dataset_path_task1=os.path.join(home_directory,dataset_path_task1)
with open(dataset_path_task1, 'rb') as handle:
whole_dataset_dicts.append(pickle.load(handle))
if 'task2' in task_name:
dataset_path_task2 = 'datasets/ZuCo/task2-NR/pickle/task2-NR-dataset.pickle'
dataset_path_task2=os.path.join(home_directory,dataset_path_task2)
with open(dataset_path_task2, 'rb') as handle:
whole_dataset_dicts.append(pickle.load(handle))
if 'task3' in task_name:
dataset_path_task3 = 'datasets/ZuCo/task3-TSR/pickle/task3-TSR-dataset.pickle'
dataset_path_task3=os.path.join(home_directory,dataset_path_task3)
with open(dataset_path_task3, 'rb') as handle:
whole_dataset_dicts.append(pickle.load(handle))
if 'taskNRv2' in task_name:
dataset_path_taskNRv2 = 'datasets/ZuCo/task2-NR-2.0/pickle/task2-NR-2.0-dataset.pickle'
dataset_path_taskNRv2=os.path.join(home_directory,dataset_path_taskNRv2)
with open(dataset_path_taskNRv2, 'rb') as handle:
whole_dataset_dicts.append(pickle.load(handle))
print()
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
# test dataset
test_set = ZuCo_dataset(whole_dataset_dicts, 'test', tokenizer, subject = subject_choice, eeg_type = eeg_type_choice, bands = bands_choice, setting = dataset_setting)
dataset_sizes = {"test_set":len(test_set)}
print('[INFO]test_set size: ', len(test_set))
# dataloaders
test_dataloader = DataLoader(test_set, batch_size = 1, shuffle=False, num_workers=4)
dataloaders = {'test':test_dataloader}
''' set up model '''
checkpoint_path = args['checkpoint_path']
pretrained_bart = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
if model_name == 'BrainTranslator':
model = BrainTranslator(pretrained_bart, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)
elif model_name == 'BrainTranslatorNaive':
model = BrainTranslatorNaive(pretrained_bart, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)
model.load_state_dict(torch.load(checkpoint_path))
model.to(device)
criterion = nn.CrossEntropyLoss()
''' eval '''
eval_model(dataloaders, device, tokenizer, criterion, model, output_all_results_path = output_all_results_path)
Datasets
Models
