Awesome
LSTM/BERT-CRF Model for Named Entity Recognition (or Sequence Labeling)
This repository implements an LSTM-CRF model for named entity recognition. The model is same as the one by Lample et al., (2016) except we do not have the last tanh
layer after the BiLSTM.
We achieve the SOTA performance on both CoNLL-2003 and OntoNotes 5.0 English datasets (check our benchmark with Glove and ELMo, other
and benchmark results with fine-tuning BERT).
Announcements
- We implemented distributed training for faster training
- We implemented a Faster CRF module which allows O(log N) inference and back-tracking!
- Benchmark results by fine-tuning BERT/Roberta**
Model | Dataset | Precision | Recall | F1 |
---|---|---|---|---|
BERT-base-cased + CRF (this repo) | CONLL-2003 | 91.69 | 92.05 | 91.87 |
Roberta-base + CRF (this repo) | CoNLL-2003 | 91.88 | 93.01 | 92.44 |
BERT-base-cased + CRF (this repo) | OntoNotes 5 | 89.57 | 89.45 | 89.51 |
Roberta-base + CRF (this repo) | OntoNotes 5 | 90.12 | 91.25 | 90.68 |
More details
Requirements
- Python >= 3.6 and PyTorch >= 1.6.0 (tested)
- pip install transformers
- pip install datasets
- pip install accelerate (optional for distributed training)
- pip install seqeval (optional, only used in evaluation while in distributed training)
In the documentation below, we present two ways for users to run the code:
- Run the model via (Fine-tuning) BERT/Roberta/etc in Transformers package.
- Run the model with simply word embeddings (and static ELMo/BERT representations loaded from external vectors).
Our default argument setup refers to the first one 1
.
Usage with Fine-Tuning BERT/Roberta (,etc) models in HuggingFace
-
Simply replace the
embedder_type
argument with the model in HuggingFace. For example, if we are usingroberta-large
, we just need to change the embedder type asroberta-large
.python transformers_trainer.py --device=cuda:0 --dataset=YourData --model_folder=saved_models --embedder_type=roberta-base
-
Distributed Training (If necessary)
- We use huggingface
accelerate
package to enable distributed training. After you set the proper configuration of your distributed environment, byaccelerate config
, you can easily run the following command for distributed training
accelerate launch transformers_trainer_ddp.py --batch_size=30 {YOUR_OTHER_ARGUMENTS}
Note that, this
batch size
is actually batch_size per GPU device. - We use huggingface
-
(Optional) Using other models in HuggingFace.
-
Run the main file with modified argument
embedder_type
:python trainer.py --embedder_type=bert-large-cased
The default value for
embedder_type
isroberta-base
. Changing the name to something likebert-base-cased
orroberta-large
, we directly load the model from huggingface. Note: if you use other models, remember to replace the tokenization mechanism inconfig/utils.py
.Our default tokenizer is assumed to be
fast_tokenizer
. If your tokenizer does not supportfast
mode, try setuse_fast=False
:tokenizer = AutoTokenizer.from_pretrained(conf.embedder_type, add_prefix_space=True, use_fast=False)
-
Finally, if you would like to know more about the details, read more details below:
- Tokenization: For BERT, we use the first wordpice to represent a complete word. Check
config/transformers_util.py
- Embedder: We show how to embed the input tokens to make word representation. Check
model/embedder/transformers_embedder.py
- Tokenization: For BERT, we use the first wordpice to represent a complete word. Check
-
Using BERT/Roberta as contextualized word embeddings (Static, Feature-based Approach) Simply go to
model/transformers_embedder.py
and uncomment the following:self.model.requires_grad = False
-
Other Usages
Using Word embedding or external contextualized embedding (ELMo/BERT) can be found in here.
Training with your own data.
- Create a folder
YourData
under the data directory. - Put the
train.txt
,dev.txt
andtest.txt
files (make sure the format is compatible, i.e. the first column is words and the last column are tags) under this directory. If you have a different format, simply modify the reader inconfig/reader.py
. - Change the
dataset
argument toYourData
when you runtrainer.py
.
Further Details and Extensions
Ongoing Plan
- Support for ELMo/BERT as features
- Interactive model where we can just import model and decode a setence
- Make the code more modularized (separate the encoder and inference layers) and readable (by adding more comments)
- Put the benchmark performance documentation to another markdown file
- Integrate BERT as a module instead of just features.
- Clean up the code to better organization (e.g.,
import
stuff) - Benchmark experiments for Transformers' based models.
- Support FP-16 training/inference
- Support distributed training using accelerate
- Releases some pre-trained NER models.
- Semi-CRF model support