Upload folder using huggingface_hub

EEG-To-Text/datasets/.gitattributes

@@ -0,0 +1,114 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.lz4 filter=lfs diff=lfs merge=lfs -text
+*.mds filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files - uncompressed
+*.pcm filter=lfs diff=lfs merge=lfs -text
+*.sam filter=lfs diff=lfs merge=lfs -text
+*.raw filter=lfs diff=lfs merge=lfs -text
+# Audio files - compressed
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+# Image files - uncompressed
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+# Image files - compressed
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+# Video files - compressed
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.webm filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZAB_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZDM_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZDN_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZGW_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZJM_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZJN_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZJS_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZKB_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZKH_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZKW_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZMG_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task1-SR/Matlab_files/resultsZPH_SR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZAB_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZDM_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZDN_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZGW_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZJM_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZJN_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZJS_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZKB_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZKH_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZKW_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZMG_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR/Matlab_files/resultsZPH_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYAC_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYAG_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYAK_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYDG_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYDR_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYFR_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYFS_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYHS_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYIS_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYLS_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYMD_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYMS_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYRH_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYRK_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYRP_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYSD_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYSL_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task2-NR-2.0/Matlab_files/resultsYTL_NR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZAB_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZDM_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZDN_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZGW_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZJM_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZJN_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZJS_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZKB_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZKH_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZKW_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZMG_TSR.mat filter=lfs diff=lfs merge=lfs -text
+ZuCo/task3-TSR/Matlab_files/resultsZPH_TSR.mat filter=lfs diff=lfs merge=lfs -text
+stanfordsentiment/stanfordSentimentTreebank/dictionary.txt filter=lfs diff=lfs merge=lfs -text

EEG-To-Text/datasets/pickle/task1-SR-datasets.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3bd2ea703a1da05b7314f17c64929a8c78307c4d2531768213bb48936702c2f
+size 1208188754

EEG-To-Text/datasets/pickle/task2-NR-2.0-datasets.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b95de456907d7851cf1807fc7bc2b65ae5242a81988ba337362b069b0b668109
+size 1758497764

EEG-To-Text/datasets/pickle/task2-NR-datasets.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb37b0f3c71dcaab4e2d49f6a1f899d2b909ef1db0adfb382f6a248e5eb8a8a9
+size 1093181866

EEG-To-Text/datasets/pickle/task3-TSR-datasets.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6dd61aa8c19bc9e98b9400edbf36cd5db4cf6434beeba0983af4bb85a632c06e
+size 1051403647

EEG-To-Text/environment.yml

@@ -0,0 +1,20 @@
+name: EEGToText
+channels:
+  - pytorch
+  - anaconda
+  - conda-forge
+  - huggingface
+dependencies:
+  - pytorch=1.9.0
+  - torchaudio=0.9.0
+  - cudatoolkit=11.1
+  - scipy=1.6.2
+  - h5py=2.10.0
+  - tqdm=4.62.0
+  - matplotlib=3.3.2
+  - transformers=4.6.1
+  - nltk=3.5
+  - pip=21.0.1
+  - pip:
+    - fuzzy-match==0.0.1
+    - rouge==1.0.0

EEG-To-Text/scripts/train_decoding.sh

@@ -0,0 +1,12 @@
+python3 train_decoding.py --model_name BrainTranslator \
+    --task_name task1_task2_task3 \
+    --one_step \
+    --pretrained \
+    --not_load_step1_checkpoint \
+    --num_epoch_step1 20 \
+    --num_epoch_step2 30 \
+    --train_input EEG \
+    -lr1 0.00002 \
+    -lr2 0.00002 \
+    -b 1 \
+    -s ./checkpoints/decoding \

EEG-To-Text/train_decoding.py

@@ -0,0 +1,379 @@
+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
+from transformers import BertLMHeadModel, BartTokenizer, BartForConditionalGeneration, BartConfig, BartForSequenceClassification, BertTokenizer, BertConfig, BertForSequenceClassification, RobertaTokenizer, RobertaForSequenceClassification, PegasusForConditionalGeneration, PegasusTokenizer, T5Tokenizer, T5ForConditionalGeneration, BertGenerationEncoder, BertGenerationDecoder, EncoderDecoderConfig, EncoderDecoderModel
+from data import ZuCo_dataset
+from model_decoding import BrainTranslator, BrainTranslatorNaive, T5Translator
+from config import get_config
+
+def train_model(dataloaders, device, model, criterion, optimizer, scheduler, num_epochs=25, checkpoint_path_best = './checkpoints/decoding/best/temp_decoding.pt', checkpoint_path_last = './checkpoints/decoding/last/temp_decoding.pt'):
+    # modified from: https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
+    since = time.time()
+      
+    best_model_wts = copy.deepcopy(model.state_dict())
+    best_loss = 100000000000
+
+    for epoch in range(num_epochs):
+        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
+        print('-' * 10)
+
+        # Each epoch has a training and validation phase
+        for phase in ['train', 'dev']:
+            if phase == 'train':
+                model.train()  # Set model to training mode
+            else:
+                model.eval()   # Set model to evaluate mode
+
+            running_loss = 0.0
+
+            # Iterate over data.
+            for input_embeddings, seq_len, input_masks, input_mask_invert, target_ids, target_mask, sentiment_labels in tqdm(dataloaders[phase]):
+                
+                # load in batch
+                input_embeddings_batch = input_embeddings.to(device).float()
+                input_masks_batch = input_masks.to(device)
+                input_mask_invert_batch = input_mask_invert.to(device)
+                target_ids_batch = target_ids.to(device)
+                """replace padding ids in target_ids with -100"""
+                target_ids_batch[target_ids_batch == tokenizer.pad_token_id] = -100
+
+                # zero the parameter gradients
+                optimizer.zero_grad()
+
+                # forward
+    	        # track history if only in train
+                with torch.set_grad_enabled(phase == 'train'):
+                    seq2seqLMoutput = model(input_embeddings_batch, input_masks_batch, input_mask_invert_batch, target_ids_batch)
+
+                    """calculate loss"""
+                    # logits = seq2seqLMoutput.logits # 8*48*50265
+                    # logits = logits.permute(0,2,1) # 8*50265*48
+
+                    # loss = criterion(logits, target_ids_batch_label) # calculate cross entropy loss only on encoded target parts
+                    # NOTE: my criterion not used
+                    loss = seq2seqLMoutput.loss # use the BART language modeling loss
+
+                    # """check prediction, instance 0 of each batch"""
+                    # print('target size:', target_ids_batch.size(), ',original logits size:', logits.size(), ',target_mask size', target_mask_batch.size())
+                    # logits = logits.permute(0,2,1)
+                    # for idx in [0]:
+                    #     print(f'-- instance {idx} --')
+                    #     # print('permuted logits size:', logits.size())
+                    #     probs = logits[idx].softmax(dim = 1)
+                    #     # print('probs size:', probs.size())
+                    #     values, predictions = probs.topk(1)
+                    #     # print('predictions before squeeze:',predictions.size())
+                    #     predictions = torch.squeeze(predictions)
+                    #     # print('predictions:',predictions)
+                    #     # print('target mask:', target_mask_batch[idx])
+                    #     # print('[DEBUG]target tokens:',tokenizer.decode(target_ids_batch_copy[idx]))
+                    #     print('[DEBUG]predicted tokens:',tokenizer.decode(predictions))
+                
+                    # backward + optimize only if in training phase
+                    if phase == 'train':
+                        # with torch.autograd.detect_anomaly():
+                        loss.sum().backward()
+                        optimizer.step()
+
+                # statistics
+                running_loss += loss.sum().item() * input_embeddings_batch.size()[0] # batch loss
+                # print('[DEBUG]loss:',loss.item())
+                # print('#################################')
+                
+
+            if phase == 'train':
+                scheduler.step()
+
+            epoch_loss = running_loss / dataset_sizes[phase]
+
+            print('{} Loss: {:.4f}'.format(phase, epoch_loss))
+
+            # deep copy the model
+            if phase == 'dev' and epoch_loss < best_loss:
+                best_loss = epoch_loss
+                best_model_wts = copy.deepcopy(model.state_dict())
+                '''save checkpoint'''
+                torch.save(model.state_dict(), checkpoint_path_best)
+                print(f'update best on dev checkpoint: {checkpoint_path_best}')
+                # with torch.set_grad_enabled(False):
+                #     traced_model_1 = torch.jit.trace(model, (torch.rand(1, 56, 840).to(device), torch.randint(1, 56).to(device), torch.rand(1, 56).to(device), torch.rand(1, 56).to(device)))
+                #     traced_model_32 = torch.jit.trace(model, (torch.rand(32, 56, 840).to(device), torch.randint(32, 56).to(device), torch.rand(32, 56).to(device), torch.rand(32, 56).to(device)))
+                # torch.jit.save(traced_model_1, checkpoint_path_best[:-3]+'_1_jit.pt')
+                # torch.jit.save(traced_model_32, checkpoint_path_best[:-3]+'_32_jit.pt')
+        print()
+
+    time_elapsed = time.time() - since
+    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
+    print('Best val loss: {:4f}'.format(best_loss))
+    torch.save(model.state_dict(), checkpoint_path_last)
+    print(f'update last checkpoint: {checkpoint_path_last}')
+
+    # load best model weights
+    model.load_state_dict(best_model_wts)
+    return model
+
+def show_require_grad_layers(model):
+    print()
+    print(' require_grad layers:')
+    # sanity check
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            print(' ', name)
+
+if __name__ == '__main__':
+    args = get_config('train_decoding')
+
+    ''' config param'''
+    dataset_setting = 'unique_sent'
+    
+    num_epochs_step1 = args['num_epoch_step1']
+    num_epochs_step2 = args['num_epoch_step2']
+    step1_lr = args['learning_rate_step1']
+    step2_lr = args['learning_rate_step2']
+    
+    batch_size = args['batch_size']
+    
+    model_name = args['model_name']
+    # model_name = 'BrainTranslatorNaive' # with no additional transformers
+    # model_name = 'BrainTranslator' 
+    
+    # task_name = 'task1'
+    # task_name = 'task1_task2'
+    # task_name = 'task1_task2_task3'
+    # task_name = 'task1_task2_taskNRv2'
+    task_name = args['task_name']
+    train_input = args['train_input']
+    print("train_input is:", train_input)   
+    save_path = args['save_path']
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+
+    skip_step_one = args['skip_step_one']
+    load_step1_checkpoint = args['load_step1_checkpoint']
+    use_random_init = args['use_random_init']
+    device_ids = [0] # device setting
+
+    if use_random_init and skip_step_one:
+        step2_lr = 5*1e-4
+        
+    print(f'[INFO]using model: {model_name}')
+    
+    if skip_step_one:
+        save_name = f'{task_name}_finetune_{model_name}_skipstep1_b{batch_size}_{num_epochs_step1}_{num_epochs_step2}_{step1_lr}_{step2_lr}_{dataset_setting}_{train_input}'
+    else:
+        save_name = f'{task_name}_finetune_{model_name}_2steptraining_b{batch_size}_{num_epochs_step1}_{num_epochs_step2}_{step1_lr}_{step2_lr}_{dataset_setting}_{train_input}'
+    
+    if use_random_init:
+        save_name = 'randinit_' + save_name
+
+    save_path_best = os.path.join(save_path, 'best')
+    if not os.path.exists(save_path_best):
+        os.makedirs(save_path_best)
+
+    output_checkpoint_name_best = os.path.join(save_path_best, f'{save_name}.pt')
+
+    save_path_last = os.path.join(save_path, 'last')
+    if not os.path.exists(save_path_last):
+        os.makedirs(save_path_last)
+
+    output_checkpoint_name_last = os.path.join(save_path_last, f'{save_name}.pt')
+
+    # subject_choice = 'ALL
+    subject_choice = args['subjects']
+    print(f'![Debug]using {subject_choice}')
+    # eeg_type_choice = 'GD
+    eeg_type_choice = args['eeg_type']
+    print(f'[INFO]eeg type {eeg_type_choice}')
+    # bands_choice = ['_t1'] 
+    # bands_choice = ['_t1','_t2','_a1','_a2','_b1','_b2','_g1','_g2'] 
+    bands_choice = args['eeg_bands']
+    print(f'[INFO]using bands {bands_choice}')
+
+
+    
+    ''' 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 = "cuda:3" 
+        dev = args['cuda'] 
+    else:  
+        dev = "cpu"
+    # CUDA_VISIBLE_DEVICES=0,1,2,3  
+    device = torch.device(dev)
+    print(f'[INFO]using device {dev}')
+    print()
+
+    ''' set up dataloader '''
+    whole_dataset_dicts = []
+    if 'task1' in task_name:
+        dataset_path_task1 = '/datasets/pickle/task1-SR-datasets.pickle'
+        with open(dataset_path_task1, 'rb') as handle:
+            whole_dataset_dicts.append(pickle.load(handle))
+    if 'task2' in task_name:
+        dataset_path_task2 = '/datasets/pickle/task2-NR-datasets.pickle'
+        with open(dataset_path_task2, 'rb') as handle:
+            whole_dataset_dicts.append(pickle.load(handle))
+    if 'task3' in task_name:
+        dataset_path_task3 = '/datasets/pickle/task3-TSR-datasets.pickle'
+        with open(dataset_path_task3, 'rb') as handle:
+            whole_dataset_dicts.append(pickle.load(handle))
+    if 'taskNRv2' in task_name:
+        dataset_path_taskNRv2 = '/datasets/pickle/task2-NR-2.0-datasets.pickle'
+        with open(dataset_path_taskNRv2, 'rb') as handle:
+            whole_dataset_dicts.append(pickle.load(handle))
+
+    print()
+
+    """save config"""
+    cfg_dir = './config/decoding/'
+
+    if not os.path.exists(cfg_dir):
+        os.makedirs(cfg_dir)
+
+    with open(os.path.join(cfg_dir,f'{save_name}.json'), 'w') as out_config:
+        json.dump(args, out_config, indent = 4)
+
+    if model_name in ['BrainTranslator','BrainTranslatorNaive']:
+        tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
+
+    elif model_name == 'PegasusTranslator':
+        tokenizer = PegasusTokenizer.from_pretrained('google/pegasus-xsum')
+    
+    elif model_name == 'T5Translator':
+        tokenizer = T5Tokenizer.from_pretrained("t5-large")
+        #tokenizer.set_prefix_tokens(language='english')
+
+    # train dataset
+    train_set = ZuCo_dataset(whole_dataset_dicts, 'train', tokenizer, subject = subject_choice, eeg_type = eeg_type_choice, bands = bands_choice, setting = dataset_setting, test_input=train_input)
+    # dev dataset
+    dev_set = ZuCo_dataset(whole_dataset_dicts, 'dev', tokenizer, subject = subject_choice, eeg_type = eeg_type_choice, bands = bands_choice, setting = dataset_setting, test_input=train_input)
+    # 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 = {'train': len(train_set), 'dev': len(dev_set)}
+    print('[INFO]train_set size: ', len(train_set))
+    print('[INFO]dev_set size: ', len(dev_set))
+    # print('[INFO]test_set size: ', len(test_set))
+    
+    # train dataloader
+    train_dataloader = DataLoader(train_set, batch_size = batch_size, shuffle=True, num_workers=4)
+    # dev dataloader
+    val_dataloader = DataLoader(dev_set, batch_size = 1, shuffle=False, num_workers=4)
+    # dataloaders
+    dataloaders = {'train':train_dataloader, 'dev':val_dataloader}
+
+    ''' set up model '''
+    if model_name == 'BrainTranslator':
+        if use_random_init:
+            config = BartConfig.from_pretrained('facebook/bart-large')
+            pretrained = BartForConditionalGeneration(config)
+        else:
+            pretrained = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
+    
+        model = BrainTranslator(pretrained, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)
+    
+    elif model_name == 'BrainTranslatorNaive':
+        pretrained = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
+        model = BrainTranslatorNaive(pretrained, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)
+
+    elif model_name == 'PegasusTranslator':
+        pretrained = PegasusForConditionalGeneration.from_pretrained('google/pegasus-xsum')
+        model = BrainTranslator(pretrained, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)
+    
+    elif model_name == 'T5Translator':
+        pretrained = T5ForConditionalGeneration.from_pretrained("t5-large")
+        model = T5Translator(pretrained, in_feature = 105*len(bands_choice), decoder_embedding_size = 1024, additional_encoder_nhead=8, additional_encoder_dim_feedforward = 2048)
+    
+    model.to(device)
+    model = torch.nn.DataParallel(model, device_ids=device_ids)
+    
+    ''' training loop '''
+
+    ######################################################
+    '''step one trainig: freeze most of BART params'''
+    ######################################################
+
+    # closely follow BART paper
+    if model_name in ['BrainTranslator','BrainTranslatorNaive', 'PegasusTranslator', 'T5Translator']:
+        for name, param in model.named_parameters():
+            if param.requires_grad and 'pretrained' in name:
+                if ('shared' in name) or ('embed_positions' in name) or ('encoder.layers.0' in name):
+                    continue
+                else:
+                    param.requires_grad = False
+
+    elif model_name == 'BertGeneration':
+        for name, param in model.named_parameters():
+            if param.requires_grad and 'pretrained' in name:
+                if ('embeddings' in name) or ('encoder.layer.0' in name):
+                    continue
+                else:
+                    param.requires_grad = False
+ 
+
+    if skip_step_one:
+        if load_step1_checkpoint:
+            stepone_checkpoint = 'path_to_step_1_checkpoint.pt'
+            print(f'skip step one, load checkpoint: {stepone_checkpoint}')
+            model.load_state_dict(torch.load(stepone_checkpoint))
+        else:
+            print('skip step one, start from scratch at step two')
+    else:
+
+        ''' set up optimizer and scheduler'''
+        optimizer_step1 = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=step1_lr, momentum=0.9)
+
+        exp_lr_scheduler_step1 = lr_scheduler.StepLR(optimizer_step1, step_size=20, gamma=0.1)
+
+        ''' set up loss function '''
+        criterion = nn.CrossEntropyLoss()
+
+        print('=== start Step1 training ... ===')
+        # print training layers
+        show_require_grad_layers(model)
+        # return best loss model from step1 training
+        model = train_model(dataloaders, device, model, criterion, optimizer_step1, exp_lr_scheduler_step1, num_epochs=num_epochs_step1, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last)
+
+    ######################################################
+    '''step two trainig: update whole model for a few iterations'''
+    ######################################################
+    for name, param in model.named_parameters():
+        param.requires_grad = True
+
+    ''' set up optimizer and scheduler'''
+    optimizer_step2 = optim.SGD(model.parameters(), lr=step2_lr, momentum=0.9)
+
+    exp_lr_scheduler_step2 = lr_scheduler.StepLR(optimizer_step2, step_size=30, gamma=0.1)
+
+    ''' set up loss function '''
+    criterion = nn.CrossEntropyLoss()
+    
+    print()
+    print('=== start Step2 training ... ===')
+    # print training layers
+    show_require_grad_layers(model)
+    
+    '''main loop'''
+    trained_model = train_model(dataloaders, device, model, criterion, optimizer_step2, exp_lr_scheduler_step2, num_epochs=num_epochs_step2, checkpoint_path_best = output_checkpoint_name_best, checkpoint_path_last = output_checkpoint_name_last)
+
+    # '''save checkpoint'''
+    # torch.save(trained_model.state_dict(), os.path.join(save_path,output_checkpoint_name))