alibaba-nlp / muver

[EMNLP 2021] MuVER: Improving First-Stage Entity Retrieval with Multi-View Entity Representations

MuVER

This repo contains the code and pre-trained model for our EMNLP 2021 paper:
MuVER: Improving First-Stage Entity Retrieval with Multi-View Entity Representations. Xinyin Ma, Yong Jiang, Nguyen Bach, Tao Wang, Zhongqiang Huang, Fei Huang, Weiming Lu

Quick Start

1. Requirements

The requirements for our code are listed in requirements.txt, install the package with the following command:

pip install -r requirements.txt

For huggingface/transformers, we tested it under version 4.1.X and 4.2.X.

2. Download data and model

• Data:
  We follow facebookresearch/BLINK to download and preprocess data. See instructions about how to download and convert to BLINK format. You will get a folder with the following structure:

- zeshel
  | - mentions
  | - documents
  | - blink_format 

• Model:
  Model for zeshel can be downloaded on https://drive.google.com/file/d/1BBTue5Vmr3MteGcse-ePqplWjccqm9_A/view?usp=sharing

3. Use the released model to reproduce our results

• Without View Merging:

export PYTHONPATH='.'  
CUDA_VISIBLE_DEVICES=YOUR_GPU_DEVICES python muver/multi_view/train.py 
    --pretrained_model path_to_model/bert-base 
    --dataset_path path_to_dataset/zeshel
    --bi_ckpt_path path_to_model/best_zeshel.bin 
    --max_cand_len 40 
    --max_seq_len 128
    --do_test 
    --test_mode test 
    --data_parallel 
    --eval_batch_size 16
    --accumulate_score

Expected Result:

World	R@1	R@2	R@4	R@8	R@16	R@32	R@50	R@64
forgotten_realms	0.6208	0.7783	0.8592	0.8983	0.9342	0.9533	0.9633	0.9700
lego	0.4904	0.6714	0.7690	0.8357	0.8791	0.9091	0.9208	0.9249
star_trek	0.4743	0.6130	0.6967	0.7606	0.8159	0.8581	0.8805	0.8919
yugioh	0.3432	0.4861	0.6040	0.7004	0.7596	0.8201	0.8512	0.8672
total	0.4496	0.5970	0.6936	0.7658	0.8187	0.8628	0.8854	0.8969

• With View Merging:

export PYTHONPATH='.'  
CUDA_VISIBLE_DEVICES=YOUR_GPU_DEVICES python muver/multi_view/train.py 
    --pretrained_model path_to_model/bert-base 
    --dataset_path path_to_dataset/zeshel
    --bi_ckpt_path path_to_model/best_zeshel.bin 
    --max_cand_len 40 
    --max_seq_len 128 
    --do_test 
    --test_mode test 
    --data_parallel 
    --eval_batch_size 16
    --accumulate_score
    --view_expansion  
    --merge_layers 4  
    --top_k 0.4

Expected result:

World	R@1	R@2	R@4	R@8	R@16	R@32	R@50	R@64
forgotten_realms	0.6175	0.7867	0.8733	0.9150	0.9375	0.9600	0.9675	0.9708
lego	0.5046	0.6889	0.7882	0.8449	0.8882	0.9183	0.9324	0.9374
star_trek	0.4810	0.6253	0.7121	0.7783	0.8271	0.8706	0.8935	0.9030
yugioh	0.3444	0.5027	0.6322	0.7300	0.7902	0.8429	0.8690	0.8826
total	0.4541	0.6109	0.7136	0.7864	0.8352	0.8777	0.8988	0.9084

Optional Argument:

• --data_parallel: whether you want to use multiple gpus.
• --accumulate_score: accumulate score for each entity. Obtain a higher score but will take much time to inference.
• --view_expansion: whether you want to merge and expand view.
• --top_k: top_k pairs are expected to merge in each layer.
• --merge_layers: the number of layers for merging.
• --test_mode: If you want to generate candidates for train/dev set, change the test_mode to train or dev, which will generate candidates outputs and save it under the directory where you save the test model.

4. How to train your MuVER

We provice the code to train your MuVER. Train the code with the following command:

export PYTHONPATH='.'  
CUDA_VISIBLE_DEVICES=YOUR_GPU_DEVICES python muver/multi_view/train.py 
    --pretrained_model path_to_model/bert-base 
    --dataset_path path_to_dataset/zeshel
    --epoch 30 
    --train_batch_size 128 
    --learning_rate 1e-5 
    --do_train --do_eval 
    --data_parallel 
    --name distributed_multi_view

Important: Since constrastive learning relies heavily on a large batch size, as reported in our paper, we use eight v100(16g) to train our model. The hyperparameters for our best model are in logs/zeshel_hyper_param.txt

The code will create a directory runtime_log to save the log, model and the hyperparameter you used. Everytime you trained your model(with or without grid search), it will create a directory under runtime_log/name_in_your_args/start_time, e.g., runtime_log/distributed_multi_view/2021-09-07-15-12-21, to store all the checkpoints, curve for visualization and the training log.

muver's Issues

关于 [train/valid/test]_top_64_candidates.tsv 文件

你好！
我想请问一下，

MuVER/muver/multi_view/data_loader.py

Line 283 in d1da6eb

candidate_path = os.path.join(candidate_path, "{}_top_64_candidates.tsv".format(mode))

这一行的 *_top_64_candidates.tsv 这个是自己处理得到的？还是zeshel原始论文里面说的使用tf-idf检索得到的64个候选？

一个bug

你好！

MuVER/muver/multi_view/train.py

Line 221 in 28761d9

if args.do_train:

这一行，应该改为：
if args.do_train and args.local_rank in [0, -1]:

按照代码逻辑，local rank不是0/1的进程 log==None，对于他们来说并没有checkpoint_path这个参数

What --pretrained_model should I use?

Hi, thanks for your contribution. When I tried to reproduce the result by using provided checkpoint, --pretrained_model arguments indicates path_to_model/bert-base, but I don't know which model should I use ( what is 'path_to_model' you used in training?). Can you help me? Thank you!

how long to test the model without muti view？why so long for me to exec it in v100（32g）？