Adding a Transformer based NER model into NERDS

In December last year at PyData LA, I did a presentation on NERDS, a toolkit fo Named Entity Recognition (NER), open sourced by some of my colleagues at Elsevier. Slides are here in case you missed it, and organizers have released the talk video as well. NERDS is a toolkit that aims to provide easy to use NER functionality for data scientists. It does so by wrapping third party NER models and exposing them through a common API, allowing data scientists to process their training data once, then train and evaluate multiple NER models (each NER model also allows for multiple tuning hyperparameters) with very little effort. In this post, I will describe a Transformer based NER model that I added recently to the 7 NER models already available my fork of NERDS.

But first, I wanted to clear up something about my talk. I was mistaken when I said that ELMo embeddings, used in Anago's ELModel and available in NERDS as ElmoNER, was subword-based, it is actually character-based. My apologies to the audience at PyData LA for misleading and many thanks to Lan Guo for catching it and setting me straight.

The Transformer architecture became popular sometime beginning of 2019, with Google's release of the BERT (Bidirectional Encoder Representations from Transformers) model. BERT was a language model that was pre-trained on large quantities of text to predict masked tokens in a text sequence, and to predict the next sentence given the previous sentence. Over the course of the year, many more BERT-like models were trained and released into the public domain, each with some critical innovation, and each performing a little better than the previous ones. These models could then be further enhanced by the user community with smaller volumes of domain specific texts to create domain-aware language models, or fine-tuned with completely different datasets for a variety of downstream NLP tasks, including NER.

The Transformers library from Hugging Face provides models for various fine-tuning tasks that can be called from your Pytorch or Tensorflow 2.x client code. Each of these models are backed by a specific Transformer language model. For example, the BERT-based fine-tuning model for NER is the BertForTokenClassification class, the structure of which is shown below. Thanks to the Transformers library, you can treat this as a tensorflow.keras.Model or a torch.nn.Module in your Tensorflow 2.x and Pytorch code respectively.

BertForTokenClassification(
  (bert): BertModel(
    (embeddings): BertEmbeddings(
      (word_embeddings): Embedding(28996, 768, padding_idx=0)
      (position_embeddings): Embedding(512, 768)
      (token_type_embeddings): Embedding(2, 768)
      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
      (dropout): Dropout(p=0.1, inplace=False)
    )
    (encoder): BertEncoder(
      (layer): ModuleList(
        (0): BertLayer(
          (attention): BertAttention(
            (self): BertSelfAttention(
              (query): Linear(in_features=768, out_features=768, bias=True)
              (key): Linear(in_features=768, out_features=768, bias=True)
              (value): Linear(in_features=768, out_features=768, bias=True)
              (dropout): Dropout(p=0.1, inplace=False)
            )
            (output): BertSelfOutput(
              (dense): Linear(in_features=768, out_features=768, bias=True)
              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
              (dropout): Dropout(p=0.1, inplace=False)
            )
          )
          (intermediate): BertIntermediate(
            (dense): Linear(in_features=768, out_features=3072, bias=True)
          )
          (output): BertOutput(
            (dense): Linear(in_features=3072, out_features=768, bias=True)
            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
        ... 11 more BertLayers (1) through (11) ...
      )
    )
    (pooler): BertPooler(
      (dense): Linear(in_features=768, out_features=768, bias=True)
      (activation): Tanh()
    )
  )
  (dropout): Dropout(p=0.1, inplace=False)
  (classifier): Linear(in_features=768, out_features=8, bias=True)
)

The figure below is from a slide in my talk, showing at a high level how fine-tuning a BERT based NER works. Note that this setup is distinct from the setup where you merely use BERT as a source of embeddings in a BiLSTM-CRF network. In a fine-tuning setup such as this, the model is essentially the BERT language model with a fully connected network attached to its head. You fine-tune this network by training it with pairs of token and tag sequences and a low learning rate. Fewer epochs of training are needed because the weights of the pre-trained BERT language model layers are already optimized and need only be updated a little to accommodate the new task.

There was also a question at the talk about whether there was a CRF involved. I didn't think there was a CRF layer at the time, but I wasn't sure, but my understanding now is that the TokenClassification models from the Hugging Face transformers library don't involve a CRF layer. This is mainly because they implement the model described in the paper BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin, Chang, Lee, and Toutanova, 2018), and that does not use a CRF. There have been some experiments such as this one, where the addition of a CRF did not seem to appreciably improve performance.

Even though using the Hugging Face transformers library is an enormous advantage compared to building this stuff up from scratch, much of the work in a typical NER pipeline is to pre-process our input into a form needed to train or predict with the fine-tuning model, and post-processing the output of the model to a form usable by the pipeline. Input to a NERDS pipeline is in the standard IOB format. A sentence is supplied as a tab separated file of tokens and corresponding IOB tags, such as that shown below:

Mr         B-PER
.          I-PER
Vinken     I-PER
is         O
chairman   O
of         O
Elsevier   B-ORG
N          I-ORG
.          I-ORG
V          I-ORG
.          I-ORG
,          O
the        O
Dutch      B-NORP
publishing O
group      O
.          O

This input gets transformed into the NERDS standard internal format (in my fork) as a list of tokenized sentences and labels:

data:   [['Mr', '.', 'Vinken', 'is', 'chairman', 'of', 'Elsevier', 'N', '.', 'V', '.', ',', 
          'the', 'Dutch', 'publishing', 'group', '.']]
labels: [['B-PER', 'I-PER', 'I-PER', 'O', 'O', 'O', 'B-ORG', 'I-ORG', 'I-ORG', 'I-ORG', 'I-ORG', 'O', 
          'O', 'B-NORP', 'O', 'O', 'O']]

Each sequence of tokens then gets tokenized by the appropriate word-piece tokenizer (in case of our BERT example, the BertTokenizer, also provided by the Transformers library). Word-piece tokenization is a way to eliminate or minimize the occurrence of unknown word lookups from the model's vocabulary. Vocabularies are finite, and in the past, if a token could not be found in the vocabulary, it would be treated as an unknown word, or UNK. Word-piece tokenization tries to match whole words as far as possible, but if it is not possible, it will try to represent a word as an aggregate of word pieces (subwords or even characters) that are present in its vocabulary. In addition (and this is specific to the BERT model, other models have different special tokens and rules about where they are placed), each sequence needs to be started using the [CLS] special token, and separated from the next sentence by the [SEP] special token. Since we only have a single sentence for our NER use case, the token sequence for the sentence is terminated with the [SEP] token. Thus, after tokenizing the data with the BertTokenizer, and applying the special tokens, the input looks like this:

[['[CLS]', 'Mr', '.', 'Vin', '##ken', 'is', 'chairman', 'of', 'El', '##se', '##vier', 'N', '.', 'V', '.', 
  ',', 'the', 'Dutch', 'publishing', 'group', '.', '[SEP]']]

This tokenized sequence will need to be featurized so it can be fed into the BertForTokenClassification network. The BertForTokenClassification only mandates the input_ids and label_ids (for training), which are basically ids for the matched tokens in the model's vocabulary and label index respectively, padded (or truncated) to the standard maximum sequence length using the [PAD] token. However, the code in run_ner.py example in the huggingface/transformers repo also builds the attention_mask (also known as masked_positions) and token_type_ids (also known as segment_ids). The former is a mechanism to avoid performing attention on [PAD] tokens, and the latter is used to distinguish between the positions for the first and second sentence. In our case, since we have a single sentence, the token_type_ids are all 0 (first sentence).

There is an additional consideration with respect to word-piece tokenization and label IDs. Consider the PER token sequence ['Mr', '.', 'Vinken'] in our example. The BertTokenizer has tokenized this to ['Mr', '.', 'Vin', '##ken']. The question is how do we distribute our label sequence ['B-PER', 'I-PER', 'I-PER']. One possibility is to ignore the '##ken' word-piece and assign it the ignore index of -100. Another possibility, suggested by Ashutosh Singh, is to treat the '##ken' token as part of the PER sequence, so the label sequence becomes ['B-PER', 'I-PER', 'I-PER', 'I-PER'] instead. I tried both approaches and did not get a significant performance bump one way or the other. Here we adopt the strategy of ignoring the '##ken' token.

Here is what the features look like for our single example sentence.

input_ids	101 1828 119 25354 6378 1110 3931 1104 2896 2217 15339 151 119 159 119 117 1103 2954 5550 1372 119 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
attention_mask	1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
token_type_ids	0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
labels	-100 5 6 6 -100 3 3 3 1 -100 -100 4 4 4 4 3 3 2 3 3 3 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100 -100

On the output side, during predictions, predictions will be generated against the input_id, attention_mask, and token_type_ids, to produce predicted label_ids. Note that the predictions are at the word-piece level and your labels are at the word level. So in addition to converting your label_ids back to actual tags, you also need to make sure that you align the prediction and label IOB tags so they are aligned.

The Transformers library provides utility code in its github repository to do many of these transformations, not only for its BertForTokenClassification model, but for its other supported Token Classification models as well. However, it does not expose the functionality through its library. As a result, your options are to either attempt to adapt the example code to your own Transformer model, or copy over the utility code into your project and import functionality from it. Because a BERT based NER was going to be only one of many NERs in NERDS, I went with the first option and concentrated only on building a BERT based NER model. You can see the code for my BertNER model. Unfortunately, I was not able to make it work well (and I think I know why as I write this post, I will update the post with my findings if I am able to make it perform better^**).

As I was building this model, adapting bits and pieces of code from the Transformers NER example code, I would often wish that they would make the functionality accessible through the library. Fortunately for me, Thilina Rajapakse, the creator of SimpleTransformers library, had the same thought. SimpleTransformers is basically an elegant wrapper on top of the Transformers library and its example code. It exposes a very simple and easy to use API to the client, and does a lot of the heavy lifting behind the scenes using the Hugging Face transformers library.

I was initially hesitant about having to add more library dependencies to NERDS (a NER based on the SimpleTransformers library needs the Hugging Face transformers library, which I had already, plus pandas and simpletransformers). However, even apart from the obvious maintainability aspect of fewer lines of code, a TransformerNER is potentially able to use all the language models supported by the underlying SimpleTransformers library - at this time, the SimpleTransformers NERModel supports BERT, RoBERTa, DistilBERT, CamemBERT, and XLM-RoBERTa language models. So adding a single TransformerNER to NERDS allows it to access 5 different Transformer Language Model backends! So the decision to switch from a standalone BertNER that relied directly on the Hugging Face transformers library, versus a TransformerNER that relied on the SimpleTransformers library was almost a no-brainer.

Here is the code for the new TransformerNER model in NERDS. As outlined in my previous blog post about Incorporating the FLair NER into NERDS, you also need to list the additional library dependencies, hook up the model so it is callable in the nerds.models package, create a short repeatable unit test, and provide some usage examples (with BioNLP, with GMB). Notice that, compared to the other NER models, we have an additional call to align the labels and predictions -- this is to correct for the word-piece tokenization creating sequences that are too long and therefore get truncated. One way around this could be to set a higher maximum_sequence_length parameter.

Performance-wise, the TransformerNER with the BERT bert-base-cased model scored the highest (average weighted F1-score) among the NERs already available in NERDS (using default hyperparameters) against both the NERDS example datasets GMB and BioNLP. The classification reports are shown below.

GMB	BioNLP
precision recall f1-score support art 0.11 0.24 0.15 97 eve 0.41 0.55 0.47 126 geo 0.90 0.88 0.89 14016 gpe 0.94 0.96 0.95 4724 nat 0.34 0.80 0.48 40 org 0.80 0.81 0.81 10669 per 0.91 0.90 0.90 10402 tim 0.89 0.93 0.91 7739 micro avg 0.87 0.88 0.88 47813 macro avg 0.66 0.76 0.69 47813 weighted avg 0.88 0.88 0.88 47813	precision recall f1-score support cell_line 0.80 0.60 0.68 1977 cell_type 0.75 0.89 0.81 4161 protein 0.88 0.81 0.84 10700 DNA 0.84 0.82 0.83 2912 RNA 0.85 0.79 0.82 325 micro avg 0.83 0.81 0.82 20075 macro avg 0.82 0.78 0.80 20075 weighted avg 0.84 0.81 0.82 20075

So anyway, really just wanted to share the news that we now have a TransformerNER model in NERDS using which you leverage what is pretty much the cutting edge in NLP technology today. I have been wanting to play with the Hugging Face transformers library for a while, and this seemed like a good opportunity initially, and the good news is that I have been able to apply this learning to simpler architectures at work (single and double sentence models using BertForSequenceClassification). However, the SimpleTransformers library from Thilina Rajapakse definitely made my job much easier -- thanks to his efforts, NERDS has an NER implementation that is at the cutting edge of NLP, and more maintainable and powerful at the same time.

^**Update (Jan 21, 2020): I had thought that the poor performance I was seeing on the BERT NER was caused by the incorrect preprocessing (I was padding first and then adding the [CLS] and [SEP] where I should have been doing the opposite), so I fixed that, and that improved it somewhat, but results are still not comparable to those from TransformerNER. I suspect it may be the training schedule in run_ner.py which is unchanged in SimpleTransformers, compared to adapted (simplified) in case of my code.

Incorporating the Flair NER into NERDS

Earlier this month I was at PyData LA where I talked about NERDS, a toolkit for Named Entity Recognition (NER) open sourced by some of my colleagues at Elsevier. You can find the slides for my talk here, the video doesn't seem to be released yet unfortunately. I covered some of this in my trip report already, but for those of you who may not know about NERDS, it is a toolkit that provides easy to use NER capabilities for data scientists. Specifically, it wraps a few (4 in the master brach, 6 in my fork -- but more on that later) third party NER models, and provides a common API for training and evaluating them. Each model also provides tunable hyperparameters and structural parameters, so as a NERDS user, you can prepare your data once and have the ability to train many different NER models quickly and efficiently.

One of the things I had promised to talk about in my abstract was how to add new NER models to NERDS, which I ended up not doing due to shortage of time. This was doubly unfortunate, because one of my aims in giving this talk was to popularize the toolkit and also to encourage contributions from Open Source to give future users of NERDS more choices. In any case, I recently added a NER from the Flair project from Zalando Research into NERDS, and figured that this might be a good opportunity to describe the steps, for the benefit of those who might be interested in extending NERDS with your own favorite third party NER model. So that's what this blog post is about.

One thing to remember though, is that, at least for now, these instructions are valid only on my fork of NERDS. In order to support the common API, NERDS exposes a common data format across all its models, and behind the scenes, converts between this format and internal formats of each model. Quite frankly, I think this is a genius idea -- an awesome application of Software Engineering principles to solve a Data Science problem. However, the common data format was somewhat baroque and a source of errors (the BiLSTM-CRF model from the Anago project on the master branch crashes intermittently because of some insidious bug which I wasn't able to crack), so I switched to a simpler data format and the bug disappeared (see the README.md for details). So we basically keep the genius idea but simplified the implementation.

Another major change is to inject parameters at construction time rather than separately during calls to fit() and predict() -- this is in line with how scikit-learn does it too, which is also where we want to go, for interoperability reasons. In any case, here is the full list of changes in the branch so far.

At a high level, here is the list of things you need to do to integrate your favorite NER into NERDS. I describe each step in greater detail below.

1. Add library dependency in setup.py
2. Figure out the third party NER API
3. Update the __init__.py file
4. Create the NERDS NER Model
5. Write and run the tests
6. Update the examples

Add library dependency in setup.py

The Flair package is installable via "pip install", so if you add it to the NERDS setup.py file as shown, it will be added to your (development) environment the next time you run "make install". The development environment simply means that the Python runtime will point to your development directory instead of somewhere central in site-packages. That way changes you make to the code will be reflected in the packag without you having to push (perhaps by additional "make install") your changes each time.

Figure out the third party NER API

If you are looking to add a NER model whose API you are already familiar with, this step may not be needed. For me, though, the Flair NER was new, so I wanted to get familiar with its API before I tried to integrate it into NERDS. I found this Flair tutorial on Tagging your Text particularly useful.

From this tutorial, I was able to figure out that Flair provides a way to train and evaluate its SequenceTagger (what we will use for our NERDS Flair NER) in one go, using a Corpus object, which is a collection of training, validation, and test datasets. Each of these datasets is a collection of Flair Sentence objects, which represents an individual sentence. Each Sentence object contains a collection of Token objects, and each Token object contains a collection of Tag objects.

Conversely, all NERDS models extends the abstract class NERModel, which inherits from the BaseEstimator and ClassifierMixin classes from scikit-learn, and expose the following four methods -- fit, predict, save, and load, as shown below. Here the fit(X, y) method is used for training the model, using dataset X and label set y. Conversely, the predict(X) method is meant for predicting labels for dataset X using a trained model. Therefore, clearly the single Corpus approach will not work for us. Luckily, however, it is possible to pass an empty Sentence list for the test dataset when creating a Corpus for training, and prediction can be done directly against the test Sentence list.

1 2 3 4 5

class NERModel(BaseEstimator, ClassifierMixin): def fit(self, X, y): pass def predict(self, X): pass def save(self, dirpath): pass def load(self, dirpath): pass

A typical train-save-load-predict pipeline consists in training the model with a labeled dataset, then saving the trained model to disk, then retrieving the saved model, and running predictions against the test set. My focus was mainly to figure out how to separate out the training and prediction code blocks into their own independent chunks, so I can reuse them in the fit() and predict(). Also, load() and save() can be somewhat idiosyncratic, with different models using different serialization mechanisms, and writing out different artifacts, so its good to watch those too. Another thing to note are the two functions sentences_to_data_labels() and data_labels_to_sentences(), that convert between the NERDS common data format (data=lists of lists of tokens, labels=lists of lists of tags), and the Sentence and Corpus based Flair data format. Its not required, of course, but I find it useful to encapsulate the conversion inside their own routines, that way they can be easily ported, not only into the final NER Model, but can potentially be reused in case I need to incorporate another NER with similar native APIs.

Here is my NER train-save-load-predict pipeline that uses the Flair NER directly. Idea is to ran this for couple of epochs just to make sure it works, and then you are ready for the next step.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93

import flair import os from flair.data import Corpus, Sentence, Token from flair.embeddings import CharacterEmbeddings, TokenEmbeddings, WordEmbeddings, StackedEmbeddings from flair.models import SequenceTagger from flair.trainers import ModelTrainer from nerds.utils import load_data_and_labels from sklearn.model_selection import train_test_split DATA_DIR = "examples/BioNLP/data" def data_labels_to_sentences(data, labels=None): sentences = [] is_dummy_labels = False if labels is None: labels = data is_dummy_labels = True for tokens, tags in zip(data, labels): sentence = Sentence() for token, tag in zip(tokens, tags): t = Token(token) if not is_dummy_labels: t.add_tag("ner", tag) sentence.add_token(t) sentences.append(sentence) return sentences def sentences_to_data_labels(sentences): data, labels = [], [] for sentence in sentences: tokens = [t.text for t in sentence.tokens] tags = [t.tags["ner"].value for t in sentence.tokens] data.append(tokens) labels.append(tags) return data, labels # training (fit) train_filename = os.path.join(DATA_DIR, "train", "Genia4ERtask1.iob2") train_data, train_labels = load_data_and_labels(train_filename) trn_data, val_data, trn_labels, val_labels = train_test_split( train_data, train_labels, test_size=0.1) trn_sentences = data_labels_to_sentences(trn_data, trn_labels) val_sentences = data_labels_to_sentences(val_data, val_labels) train_corpus = Corpus(trn_sentences, val_sentences, [], name="train-corpus") print(train_corpus) basedir = "flair-ner-test" savedir = "flair-saved" tag_dict = train_corpus.make_tag_dictionary(tag_type="ner") embedding_types = [ WordEmbeddings("glove"), CharacterEmbeddings() ] embeddings = StackedEmbeddings(embeddings=embedding_types) tagger = SequenceTagger(hidden_size=256, embeddings=embeddings, tag_dictionary=tag_dict, tag_type="ner", use_crf=True) trainer = ModelTrainer(tagger, train_corpus) trainer.train(basedir, learning_rate=0.1, mini_batch_size=32, max_epochs=2) # model is saved by default, but let's do it again os.makedirs(savedir, exist_ok=True) tagger.save(os.path.join(savedir, "final-model.pt")) # load back the model we trained model_r = SequenceTagger.load(os.path.join(savedir, "final-model.pt")) # prediction (predict) test_filename = os.path.join(DATA_DIR, "test", "Genia4EReval1.iob2") test_data, test_labels = load_data_and_labels(test_filename) test_sentences = data_labels_to_sentences(test_data) pred_sentences = model_r.predict(test_sentences, mini_batch_size=32, all_tag_prob=True) i = 0 _, predictions = sentences_to_data_labels(pred_sentences) for prediction in predictions: print(prediction) i += 1 if i > 10: break

The resulting model is shown below. It looks similar to the word+character hybrid model proposed by Guillaume Genthial in his Sequence Tagging with Tensorflow blog post, where word embeddings (seeded with GloVe vectors) and embeddings generated from characters are concatenated and fed into an LSTM, and then the output of the LSTM is fed into a linear layer with CRF loss to produce the predictions.

1 2 3 4 5 6 7 8 9 10 11 12 13 14

SequenceTagger( (embeddings): StackedEmbeddings( (list_embedding_0): WordEmbeddings('glove') (list_embedding_1): CharacterEmbeddings( (char_embedding): Embedding(275, 25) (char_rnn): LSTM(25, 25, bidirectional=True) ) ) (word_dropout): WordDropout(p=0.05) (locked_dropout): LockedDropout(p=0.5) (embedding2nn): Linear(in_features=150, out_features=150, bias=True) (rnn): LSTM(150, 256, batch_first=True, bidirectional=True) (linear): Linear(in_features=512, out_features=20, bias=True) )

Update the __init__.py file

Python's package paths are very file-oriented. For example, functions in the nerds.utils package are defined in the nerds/utils.py file. However, since NER models are typically large blocks of code, my preference (as well as the original authors) is to have each model in its own file. This can lead to very deep package structures, or we can effectively flatten the package paths by importing them into the nerds.models package in the nerds/models/__init__.py. You can now refer to the FlairNER class defined in nerds/models/flair.py as nerds.models.FlairNER.

Create the NERDS NER model

At this point, it is fairly easy to build the FlairNER class with code chunks from the throwaway train-save-load-predict script. There are a few things to keep in mind that have to do in part with coding style, and in part with a desire for interoperability with scikit-learn and its huge library of support functions. I try to follow the guidelines in Developing scikit-learn estimators. One important deviation from the guidelines is that we don't allow **kwargs for fit() and predict(), since its easier to track the parameters if they are all passed in via the constructor. Another important thing to note is that NERDS models are not true Estimators, since fit and predict work with lists of lists of primitive objects, rather than just lists, so the check_estimator function fails on these models -- although I think this may be because the creators of check_estimator may not have anticipated this usage.

We don't have publicly available API Docs for NERDS yet, but in anticipation of that, we are using the NumPy DocString format as our guide, as advised by the Scikit-Learn coding guidelines.

Finally, in the save() function, we dump out the parameters fed into the constructor in a YAML file. This is mainly for documentation purposes, to save the user time figuring out after the fact which model was created with which hyperparameters. The class structure doesn't enforce this requirement, i.e., the NER will happily work even without this feature, but its a single-line call to utils.write_param_file(), so its not a lot of work for something very useful, so you just have to remember to add this in.

Here is the code for the FlairNER class. As you can see, a lot of code has been copy-pasted from the throwaway train-save-load-predict code that we built earlier. There is also some validation code, for example, to prevent predict() being run without a trained model, or to complain if the code is asked to load the model from a non-existent location, etc. Also the private functions _convert_to_flar() and _convert_from_flair() are basically clones of the data_labels_to_sentences() and sentence_to_data_labels() functions from the earlier script.

Write and run the tests

NERDS has a suite of unit tests in the nerds/test directory. It uses the nose package for running the tests. For the NER Models, we have a tiny dataset of 2 sentences, with which we train and predict. The dataset is generally insufficient to train an NER model, so basically all we are looking for is that the code runs end-to-end without complaining about size issues, etc. Here is the code for the FlairNER tests.

You can run the test individually by "nosetests nerds/tests/test_flair_ner.py" or run all tests using "make test". I like to start with running individual tests to make sure my changes are good, and then follow it up with a final "make test" to make sure my changes haven't broken something elsewhere in the system.

Update the examples

Finally, it is time to add your NER to the example code in nerds/examples. This is mainly for NERDS users, to provide them examples on how to call NERDS, but it can also be interesting for you, to see how your new NER stacks up against the ones that are there already. There are two examples, one based on the Groningen Meaning Bank (GMB) dataset of general entities such as PERson, LOCation, etc., and another based on the BioNLP dataset for Bio-Entity recognition. As mentioned earlier, NERDS allows you to prepare your data once and reuse it across multiple models, so the code to include the FlairNER is this block here and here respectively. As can be seen from the classification reports on the respective README.md (here and here), performance of the FlairNER is on par with the BiLSTM-CRF in case of GMB but closer to CRF in case of BioNLP.

That's basically all it takes code-wise, to add a new NER to NERDS. The next step is of course to do a Pull Request (PR), which I would request you to hold off on at the moment, since I am working off a fork myself, and my git-fu is not powerful enough to figure how to handle PRs against a fork. I would prefer that my fork gets pulled into master first, then we handle any additional PRs. However, please queue them up on the NERDS Issues page, so they can be incorporated as they come in.