Interactive Machine Comprehension with Dynamic Knowledge Graphs

Implementation for the EMNLP 2021 paper.


apt-get -y update
apt-get install -y unzip zip parallel
conda create -p /tmp/imrc python=3.6 numpy scipy cython nltk
conda activate /tmp/imrc
pip install --upgrade pip
pip install numpy==1.16.2
pip install gym==0.15.4
pip install tqdm pipreqs pyyaml pytz visdom
conda install pytorch torchvision cudatoolkit=9.2 -c pytorch
pip install transformers
pip install allennlp

Data Preparation

Split SQuAD 1.1 and preprocess

The original SQuAD dataset does not provide its test set, we take 23 wiki articles from its training set as our validation set. We then use the SQuAD dev set as our test set.

# download SQuAD from official website, then
python utils/

To speed up training, we parse (tokenization and SRL) the dataset in advance.

python utils/

This will result squad_split/processed_squad.1.1.split.[train/valid/test].json, which are used in iMRC tasks.

Preprocess Wikipedia data for self-supervised learning

python utils/
python utils/

This will result wiki_without_squad/wiki_without_squad_[train/valid/test].json, which are used to pre-train the continuous belief graph generator.


To train the agent equipped with different types of graphs, run:

# without graph
python configs/imrc_none.yaml

# co-occurrence graph
python configs/imrc_cooccur.yaml

# relative position graph
python configs/imrc_rel_pos.yaml

# SRL graph
python configs/imrc_srl.yaml

# continuous belief graph
# in this setting, we need a pre-trained graph generator.
# we provide our pre-trained graph generator at

# one can choose to train their own version by:
python configs/pretrain_cont_bnelief.yaml
# then using the downloaded/saved model checkpoint
python configs/imrc_cont_belief.yaml

To change the task settings/configurations:

  naozi_capacity: 1  # capacity of agent's external memory queue (1, 3, 5)
  generate_or_point: "point"  # "qmpoint": q+o_t, "point": q, "generate": vocab
  disable_prev_next: False  # False: Easy Mode, True: Hard Mode

  recurrent: True  # recurrent component described in Section 3.3 and Section 4.Additional Results


  title={Interactive Machine Comprehension with Dynamic Knowledge Graphs},
  author={Xingdi Yuan},