bert-crf-entity-extraction-pytorch

This repository is for the entity extraction task using the pre-trained BERT[1] and the additional CRF(Conditional Random Field)[2] layer.

Originally, this project has been conducted for dialogue datasets, so it contains both single-turn setting and multi-turn setting.

The single-turn setting is the same as the basic entity extraction task, but the multi-turn one is a little bit different since it considers the dialogue contexts(previous histories) to conduct the entity extraction task to current utterance.

The multi-turn context application is based on ReCoSa(the Relevant Contexts with Self-attention)[3] structure.

You can see the details of each model in below descriptions.

Arguments

Arguments for data pre-processing

Arguments for training/evaluating

Dataset

This repository uses the Google's Taskmaster-2[4] dataset for entity extraction task.

You should first download the data ("TM-2-2020"), and get all json files in "TM-2-2020/data" directory to properly run this project.

You can see the detailes for using the Taskmaster-2 dataset in the next section.

How to run

1. Install all required packages.

   pip install -r requirements.txt
   

   
   

2. Make the directory {data_dir}/{raw_dir} and put the json files, as mentioned in the previous section.

   In default setting, the structure of whole data directory should be like below.

   data
   └--raw
       └--flight.json
       └--food-ordering.json
       └--hotels.json
       └--movies.json
       └--music.json
       └--restaurant-search.json
       └--sports.json
   

   
   

3. Run the data processing script.

   sh exec_data_processing.sh
   

   After running it, you will get the processed files like below in the default setting.

   data
   └--raw
       └--flight.json
       └--food-ordering.json
       └--hotels.json
       └--movies.json
       └--music.json
       └--restaurant-search.json
       └--sports.json
   └--processed
       └--class_dict.json
       └--train_tokens.pkl
       └--train_tags.pkl
       └--valid_tokens.pkl
       └--valid_tags.pkl
       └--test_tokens.pkl
       └--test_tags.pkl
   

   
   

4. Run the main script and check the results.

   sh exec_main.sh
   

Results

References

[1] Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2018). Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. (https://arxiv.org/abs/1810.04805)

[2] Lafferty, J., McCallum, A., & Pereira, F. C. (2001). Conditional random fields: Probabilistic models for segmenting and labeling sequence data. (https://repository.upenn.edu/cis_papers/159/)

[3] Zhang, H., Lan, Y., Pang, L., Guo, J., & Cheng, X. (2019). Recosa: Detecting the relevant contexts with self-attention for multi-turn dialogue generation. arXiv preprint arXiv:1907.05339. (https://arxiv.org/abs/1907.05339)

[4] Taskmaster-2 . (2020). (https://research.google/tools/datasets/taskmaster-2/)