A PyTorch implementation of the BI-LSTM-CRF model
Project description
A PyTorch implementation of the BI-LSTM-CRF model.
Features:
- Compared with PyTorch BI-LSTM-CRF tutorial, following improvements are performed:
- Full support for mini-batch computation
- Full vectorized implementation. Specially, removing all loops in "score sentence" algorithm, which dramatically improve training performance
- CUDA supported
- Very simple APIs for CRF module
- START/STOP tags are automatically added in CRF
- A inner Linear Layer is included which transform from feature space to tag space
- Specialized for NLP sequence tagging tasks
- Easy to train your own sequence tagging models
- MIT License
Installation
- dependencies
- Python 3
- PyTorch
- install
$ pip install bi-lstm-crf
Training
corpus
- prepare your corpus in the specified structure and format
- there is also a sample corpus in
bi_lstm_crf/app/sample_corpus
training
$ python -m bi_lstm_crf corpus_dir --model_dir "model_xxx"
training curve
import pandas as pd import matplotlib.pyplot as plt # the training losses are saved in the model_dir df = pd.read_csv(".../model_dir/loss.csv") df[["train_loss", "val_loss"]].ffill().plot(grid=True) plt.show()
Prediction
from bi_lstm_crf.app import WordsTagger model = WordsTagger(model_dir="xxx") tags, sequences = model(["市领导到成都..."]) # CHAR-based model print(tags) # [["B", "B", "I", "B", "B-LOC", "I-LOC", "I-LOC", "I-LOC", "I-LOC", "B", "I", "B", "I"]] print(sequences) # [['市', '领导', '到', ('成都', 'LOC'), ...]] # model([["市", "领导", "到", "成都", ...]]) # WORD-based model
CRF Module
The CRF module can be easily embeded into other models:
from bi_lstm_crf import CRF # a BERT-CRF model for sequence tagging class BertCrf(nn.Module): def __init__(self, ...): ... self.bert = BERT(...) self.crf = CRF(in_features, num_tags) def loss(self, xs, tags): features, = self.bert(xs) masks = xs.gt(0) loss = self.crf.loss(features, tags, masks) return loss def forward(self, xs): features, = self.bert(xs) masks = xs.gt(0) scores, tag_seq = self.crf(features, masks) return scores, tag_seq
References
- Zhiheng Huang, Wei Xu, and Kai Yu. 2015. Bidirectional LSTM-CRF Models for Sequence Tagging. arXiv:1508.01991.
- PyTorch tutorial ADVANCED: MAKING DYNAMIC DECISIONS AND THE BI-LSTM CRF
Project details
Release history Release notifications
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Filename, size | File type | Python version | Upload date | Hashes |
---|---|---|---|---|
Filename, size bi-lstm-crf-0.1.2.tar.gz (9.8 kB) | File type Source | Python version None | Upload date | Hashes View hashes |