Please find more details of this work in our paper.
Check the README in the github. princeton-nlp/PURE
The input data format of the entity model is JSONL. Each line of the input file contains one document in the following format.
{
# document ID (please make sure doc_key can be used to identify a certain document)
"doc_key": "CNN_ENG_20030306_083604.6",
# sentences in the document, each sentence is a list of tokens
"sentences": [
[...],
[...],
["tens", "of", "thousands", "of", "college", ...],
...
],
# entities (boundaries and entity type) in each sentence
"ner": [
[...],
[...],
[[26, 26, "LOC"], [14, 14, "PER"], ...], #the boundary positions are indexed in the document level
...,
],
# relations (two spans and relation type) in each sentence
"relations": [
[...],
[...],
[[14, 14, 10, 10, "ORG-AFF"], [14, 14, 12, 13, "ORG-AFF"], ...],
...
]
}
You can use run_entity.py
with --do_train
to train an entity model and with --do_eval
to evaluate an entity model.
A trianing command template is as follow:
python run_entity.py \
--do_train --do_eval [--eval_test] \
--learning_rate=1e-5 --task_learning_rate=5e-4 \
--train_batch_size=16 \
--context_window {0 | 100 | 300} \
--task {ace05 | ace04 | scierc} \
--data_dir {directory of preprocessed dataset} \
--model {bert-base-uncased | albert-xxlarge-v1 | allenai/scibert_scivocab_uncased} \
--output_dir {directory of output files}
Arguments:
--learning_rate
: the learning rate for BERT encoder parameters.--task_learning_rate
: the learning rate for task-specific parameters, i.e., the classifier head after the encoder.--context_window
: the context window size used in the model.0
means using no contexts. In our cross-sentence entity experiments, we use--context_window 300
for BERT models and SciBERT models and use--context_window 100
for ALBERT models.--model
: the base transformer model. We usebert-base-uncased
andalbert-xxlarge-v1
for ACE04/ACE05 and useallenai/scibert_scivocab_uncased
for SciERC.--eval_test
: whether evaluate on the test set or not.
The predictions of the entity model will be saved as a file (ent_pred_dev.json
) in the output_dir
directory. If you set --eval_test
, the predictions (ent_pred_test.json
) are on the test set. The prediction file of the entity model will be the input file of the relation model.
The input data format of the relation model is almost the same as that of the entity model, except that there is one more filed ."predicted_ner"
to store the predictions of the entity model.
{
"doc_key": "CNN_ENG_20030306_083604.6",
"sentences": [...],
"ner": [...],
"relations": [...],
"predicted_ner": [
[...],
[...],
[[26, 26, "LOC"], [14, 15, "PER"], ...],
...
]
}
You can use run_relation.py
with --do_train
to train a relation model and with --do_eval
to evaluate a relation model. A trianing command template is as follow:
python run_relation.py \
--task {ace05 | ace04 | scierc} \
--do_train --train_file {path to the training json file of the dataset} \
--do_eval [--eval_test] [--eval_with_gold] \
--model {bert-base-uncased | albert-xxlarge-v1 | allenai/scibert_scivocab_uncased} \
--do_lower_case \
--train_batch_size 32 \
--eval_batch_size 32 \
--learning_rate 2e-5 \
--num_train_epochs 10 \
--context_window {0 | 100} \
--max_seq_length {128 | 228} \
--entity_output_dir {path to output files of the entity model} \
--output_dir {directory of output files}
Aruguments:
--eval_with_gold
: whether evaluate the model with the gold entities provided.--entity_output_dir
: the output directory of the entity model. The prediction files (ent_pred_dev.json
orent_pred_test.json
) of the entity model should be in this directory.
The prediction results will be stored in the file predictions.json
in the folder output_dir
, and the format will be almost the same with the output file from the entity model, except that there is one more field "predicted_relations"
for each document.
You can run the evaluation script to output the end-to-end performance (Ent
, Rel
, and Rel+
) of the predictions.
python run_eval.py --prediction_file {path to output_dir}/predictions.json
Performance of pretrained models on AGAC test set:
- BERT-base-cased (cross)
NER-All instance - P: 0.41667, R: 0.04677, F1: 0.08410 -- 2021-04-20
If you use this code in your research, please cite fllow work:
@inproceedings{zhong2021frustratingly,
title={A Frustratingly Easy Approach for Entity and Relation Extraction},
author={Zhong, Zexuan and Chen, Danqi},
booktitle={North American Association for Computational Linguistics (NAACL)},
year={2021}
}