MikeWangWZHL / EEG-To-Text

code for AAAI2022 paper "Open Vocabulary Electroencephalography-To-Text Decoding and Zero-shot Sentiment Classification"
162 stars 35 forks source link
bart bert brain-computer-interface eeg-classification eeg-signals language language-model sentiment-analysis

IMPORTANT UPDATE (6/29/2024)

The original code has an unintential bug in evaluation that causes an inaccurate assessment of the model’s true capability. We thank Yiqian Yang and Hyejeong Jo, et al for spotting and investigating the problem! Please read https://arxiv.org/pdf/2405.06459 for more details.

Please refer to https://github.com/NeuSpeech/EEG-To-Text for the corrected code and detailed experiments. To avoid further confusion, we archived this repo.


Please read https://arxiv.org/pdf/2405.06459 before proceeding!

(AAAI 2022) Open Vocabulary EEG-To-Text Decoding and Zero-shot sentiment classification

Create Environment

run conda env create -f environment.yml to create the conda environment (named "EEGToText") used in our experiments.

Download ZuCo datasets

Preprocess datasets

run bash ./scripts/prepare_dataset.sh to preprocess .mat files and prepare sentiment labels.

For each task, all .mat files will be converted into one .pickle file stored in ~/datasets/ZuCo/<task_name>/<task_name>-dataset.pickle.

Sentiment dataset for ZuCo (sentiment_labels.json) will be stored in ~/datasets/ZuCo/task1-SR/sentiment_labels/sentiment_labels.json.

Sentiment dataset for filtered Stanford Sentiment Treebank will be stored in ~/datasets/stanfordsentiment/ternary_dataset.json

Usage Example

Open vocabulary EEG-To-Text Decoding

To train an EEG-To-Text decoding model, run bash ./scripts/train_decoding.sh.

To evaluate the trained EEG-To-Text decoding model from above, run bash ./scripts/eval_decoding.sh.

For detailed configuration of the available arguments, please refer to function get_config(case = 'train_decoding') in /config.py

Zero-shot sentiment classification pipeline

We first train the decoder and the classifier individually, and then we evaluate the pipeline on ZuCo task1-SR data.

To run the whole training and evaluation process, run bash ./scripts/train_eval_zeroshot_pipeline.sh.

For detailed configuration of the available arguments, please refer to function get_config(case = 'eval_sentiment') in /config.py

Citation

@inproceedings{wang2022open,
  title={Open vocabulary electroencephalography-to-text decoding and zero-shot sentiment classification},
  author={Wang, Zhenhailong and Ji, Heng},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  volume={36},
  number={5},
  pages={5350--5358},
  year={2022}
}