Knowledge Graph Aligned Entity Linker using SentenceTransformers

Most of us are familiar with Named Entity Recognizers (NERs) that can recognize spans in text as belonging to a small number of classes, such as Person (PER), Organization (ORG), Location (LOC), etc. These are usually multi-class classifier models, trained on input sequences to return BIO (Begin-Input-Output) tags for each token. However, recognizing entities in a Knowledge Graph (KG) using this approach is usually a much harder proposition, since a KG can contain thousands, even millions, of distinct entities, and it is just not practical to create a multi-class classifier for so many target classes. A common approach to building a NER for such a large number of entities is to use dictionary based matching. However, the approach suffers from the inability to do "fuzzy" or inexact matching, beyond standard normalization streategies such as lowercasing and stemming / lemmatizing, and requires you to specify up-front all possible synonyms that may be used to refer to a given entity.

An alternative approach may be to train another model, called a Named Entity Linker (NEL) that would take the spans recognized as candidate entities or phrases by the NER model, and then attempt to link the phrase to an entity in the KG. In this situation, the NER just learns to predict candidate phrases that may be entities of interest, which puts it on par with simpler PER/ORG/LOC style NERs in terms of complexity. The NER and NEL are pipelined together in a setup that is usually known as Named Entity Recognition and Linking (NERL).

In this post, I will describe a NEL model that I built for my 2023 Dev10 project. Our Dev10 program allows employees to use up to 10 working days per year to pursue a side-project, similar to Google's 20% program. The objective is to learn an embedding model where encodings of synonyms of a given entity are close together, and where encodings of synonyms of different entities are pushed far apart. We can then encode each entity in this space as the encoding of the centroid of the encodings of its individual synonyms. Each candidate phrase output from the NER model can then be encoded using this embedding model, and its nearest neighbors in the embedding space would correspond to the most likely entities to link to.

The idea is inspired by Self-Alignment Pretraining for Biomedical Entity Representations (Liu et al, 2021) which produced the SapBERT model (SAP == Self Aligned Pretraining). It uses Contrastive Learning to fine-tune the BiomedBERT model. In this scenario, positive pairs are pairs of synonyms for the same entity in the KG and negative pairs are synonyms from different entities. It uses the Unified Medical Language System (UMLS) as its KG, to source synonym pairs.

I follow a largely similar approach in my project, except that I use the SentenceTransformers library to fine tune the BiomedBERT model. For my initial experiments, I also used the UMLS as my source of synonym pairs, mainly for reproducibility purposes since it is a free resource available for download to anyone. I tried fine-tuning a bert-base-uncased model and the BiomedBERT models, with MultipleNegativesRanking (MNR) as well as Triplet loss, the latter with Hard Negative Mining. My findings are in line with the SapBERT paper, i.e. that BiomedBERT performs better than BERT base, and that MNR performs better than Triplet loss. The last bit was something of a dissapointment, since I had expected Triplet loss to perform better. It is possible that the Hard Negative Mining was not hard enough, or maybe I needed a higher number than 5 negatives for each positive.

You can learn more about the project in my GitHub repository sujitpal/kg-aligned-entity-linker, as well as find the code in there, in case you want to replicate it.

Here are some visualizations from my best model. The chart on the left shows the distribution of cosine similarities between known negative synonym pairs (orange curve) and known positive synonym pairs (blue curve). As you can see, there is almost no overlap. The heatmap on the right shows the cosine similarity of a set of 10 synonym pairs, where the diagonal corresponds to positive pairs and the non-diagonal elements correspond to negative pairs. As you can see, the distribution seems quite good.

I also built a small demo that shows what in my opinion is the main use case for this model. It is a NERL pipeline, where the NER component is the UMLS entity finder (en_core_sci_sm) from the SciSpacy project, and the NEL component is my best performing model (kgnel-bmbert-mnr). In order to look up nearest neighbors for a given phrase encoding, the NEL component also needs a vector store to store the centroids of the encodings of entity synonyms, I used QDrant for this purpose. The QDrant vector store needs to be populated with the centroid embeddings in advance, and in order to cut down on the index and vectorization time, I only computed embeddings for centroids for entities of type "Disease or Syndrome" and "Clinical Drug". The visualizations below show the outputs (from displacy) of the outputs of the NER component:

and that of the NEL component in my demo NERL pipeline. Note that only spans that were identified as a Disease or Drug with a confidence above a threshold were selected in this phase.

Such a NERL pipeline could be used to mine new literature for new synonyms of existing entities. Once discovered, they could be added to the synonym list for the dictionary based NER to increase its recall.

Anyway, that was all I had for this post. Today is also January 1 2024, so I wanted to wish you all a very Happy New Year and a productive 2024 filled with many Machine Learning adventures!

Entities from CORD-19 using Dask, SciSpaCy, and Saturn Cloud

Its been a while since I last posted here, but I recently posted on our Elsevier Labs blog, and I wanted to point folks here to that. The post, titled How Elsevier Accelerated COVID-19 research using Dask and Saturn Cloud, describes some work I did to extract biomedical entities from the CORD-19 dataset using Dask and trained Named Entity Recognition (NER) and Named Entity Recognition and Linking (NERL) models from SciSpaCy, on the Saturn Cloud platform.

At a high level, the pipeline takes documents from the CORD-19 dataset as input, decomposes them into sentences, and passes each sentence through one of nine trained SciSpaCy models (4 NER and 5 NERL) to extract spans of text representing different kinds of biomedical entities, such as CHEMICAL, DISEASE, GENE, PROTEIN, etc., as well as entities listed in various well-known biomedical ontologies (such as UMLS, MeSH, etc). The output is provided in tabular format as Parquet files consumable from many platforms, including Dask and Spark.

The pipeline described in the post was developed and executed on the Saturn Cloud platform. Saturn Cloud is a Platform as a Service (PaaS) that provides a Jupyter Notebooks (or Jupyter Labs) development environment on top of Amazon Web Services (AWS). It also provides a custom Dask scheduler that allows you to scale out to a cluster of workers. It also provides RAPIDS on GPU boxes for vertical scaling (scaling up), but I didn't use RAPIDS for this work.

Before I started working with Saturn Cloud, I was trying to develop the same pipeline (also using Dask) on a single AWS EC2 box (a t2.2xlarge with 8 vCPUs and 32 GB RAM. However, after the first few steps, I rapidly began to hit the resource constraints of a single machine, leading me to some interesting workarounds I descibe here. Once I moved to Saturn Cloud, these problems largely went away because I could now scale out the processing across a cluster of machines. In addition, the code got simpler because I no longer needed to work around resource constraints imposed by my single machine environment. My Saturn Cloud notebooks are available at my Github repository sujitpal/saturn-scipacy under an Apache 2.0 license. The README.md provides additional details about how the notebooks are organized, and the format of the output files.

We built two pipelines, the first is what we call a "full" pipeline. Input to this is a dated CORD-19 dataset and it extracts different entities using the nine models and outputs the entities as structured files. The second is an incremental format, which takes entities from a previous version of the COVID-19 dataset and the current dataset, figures out document additions and deletions between the two, and generates entities only for the added documents and deletes entities corresponding to the deleted documents. The incremental pipeline completes much faster than the full pipeline, typically under an hour compared to about 15 hours.

More interesting than the details of the processing, however, is the fact that we have made the output freely available at this requester-pays bucket on AWS. You will need an AWS account to access the data. The 2020-08-28 folder represents entities extracted by the full pipeline from the September 28 2020 version of CORD-19, and the 2020-09-28 folder represents entities extracted by the incremental pipeline from the October 28 2020 version. Each dataset is about 30-35 GB in size.

Because these are relatively large datasets, it is generally advisable to bring the code to the data rather than the other way round. So you probably want to keep the data in the cloud, maybe even within AWS (in which case you won't need to pay any network charges).

We believe the entity data would be useful for NLP based biomedical work (in-silico biology and pharma). Since the input to the pipeline, as well as the models, were both in the public domain, we thought it was only fitting that the output of the pipeline also be in the public domain. We hope it helps to advance the state of scientific knowledge around COVID-19 and helps in humanity's fight against the pandemic. If you happen to use the data in an academic setting, we would appreciate you citing it as Pal, Sujit (2020), “CORD-19 SciSpaCy Entity Dataset”, Mendeley Data, V2, doi: 10.17632/gk9njn3pth.2.

Disambiguating SciSpacy + UMLS entities using the Viterbi algorithm

The SciSpacy project from AllenAI provides a language model trained on biomedical text, which can be used for Named Entity Recognition (NER) of biomedical entities using the standard SpaCy API. Unlike the entities found using SpaCy's language models (at least the English one), where entities have types such as PER, GEO, ORG, etc., SciSpacy entities have the single type ENTITY. In order to further classify them, SciSpacy provides Entity Linking (NEL) functionality through its integration with various ontology providers, such as the Unified Medical Language System (UMLS), Medical Subject Headings (MeSH), RxNorm, Gene Ontology (GO), and Human Phenotype Ontology (HPO). 

The NER and NEL processes are decoupled. The NER process finds candidate entity spans, and these spans are matched against the respective ontologies, which may result in the span matching zero or more ontology entries. All candidate span is then matched to all the matched entities. 

I tried annotating the COVID-19 Open Research Dataset (CORD-19) against UMLS using the SciSpacy integration described above, and I noticed significant ambiguity in the linking results. Specifically, annotating approximately 22 million sentences in the CORD-19 dataset results in 113 million candidate entity spans, which get linked to 166 million UMLS concepts, i.e., on average, each candidate span resolves to 1.5 UMLS concepts. However, the distribution is Zipfian, with approximately 46.87% entity spans resolving to a single concept, with a long tail of entity spans being linked to up to 67 UMLS concepts. 

In this post, I will describe a strategy to disambiguate the linked entities. Based on limited testing, this chooses the correct concept about 73% of the time. 

The strategy is based on the intuition that an ambiguously linked entity span is more likely to resolve to a concept that is closely related to concepts for the other non-ambiguously linked entity spans in the sentence. In other words, the best target label to choose for an ambiguous entity is the one that is semantically closest to the labels of other entities in the sentence. Or even more succintly, and with apologies to John Firth, an entity is known by the company it keeps. 

The NER and NEL processes provided by the SciSpacy library allows us to reduce a sentence to a collection of entity spans, each of which map to zero or more UMLS concepts. Each UMLS concept maps to one or more Semantic Types, which represent high level subject categories. So essentially, a sentence can be reduced to a graph of semantic type using the following steps. 

Consider the sentence below, the NER step identifies candidate spans that are indicated by highlights.

The fact that viral antigens could not be demonstrated with the used staining is not the result of antibodies present in the cat that already bound to these antigens and hinder binding of other antibodies.

The NEL step will attempt to match these spans against the UMLS ontology. Results for the matching are shown below. As noted earlier, each UMLS concept maps to one or more sematic types, and these are shown here as well.

   

Entity-ID	Entity Span	Concept-ID	Concept Primary Name	Semantic Type Code	Semantic Type Name
1	staining	C0487602	Staining method	T059	Laboratory Procedure
2	antibodies	C0003241	Antibodies	T116	Amino Acid, Peptide, or Protein
				T129	Immunologic Factor
3	cat	C0007450	Felis catus	T015	Mammal
		C0008169	Chloramphenicol O-Acetyltransferase	T116	Amino Acid, Peptide, or Protein
				T126	Enzyme
		C0325089	Family Felidae	T015	Mammal
		C1366498	Chloramphenicol Acetyl Transferase Gene	T028	Gene or Genome
4	antigens	C0003320	Antigens	T129	Immunologic Factor
5	binding	C1145667	Binding action	T052	Activity
		C1167622	Binding (Molecular Function)	T044	Molecular Function
6	antibodies	C0003241	Antibodies	T116	Amino Acid, Peptide, or Protein
				T129	Immunologic Factor

The sequence of entity spans, each mapped to one or more semantic type codes can be represented by a graph of semantic type nodes as shown below. Here, each vertical grouping corresponds to an entity position. The BOS node is a special node representing the beginning of the sequence. Based on our intuition above, entity disambiguation is now just a matter of finding the most likely path through the graph.

Of course, "most likely" implies that we need to know the probabilities for transitioning between semantic types. We can think of the graph as a Markov Chain, and consider the probability of each node in the graph as being determined only by its previous node. Fortunately, this information is already available as a result of the NER + NEL process for the entire CORD-19 dataset, where approximately half of the entity spans mapped unambiguously to a single UMLS concept. Most concepts map to a single semantic type, but in cases where they map to multiple, we consider them as separate records. We compute pairwise transition probabilities across semantic types for these unambiguously linked pairs across the CORD-19 dataset and create our transition matrix. In addition, we also create a matrix of emission probabilities that identify the probabilities of resolving to a concept given a semantic type. 

Using the transition probabilities, we can traverse each path in the graph from starting to ending position, computing the path probability as the product of transition probabilities (or for computational reasons, the sum of log-probabilities) of the edges. However, better methods exist, such as the Viterbi algorithm, which allows us to save on repeated computation of common edge sequences across multiple paths. This is what we used to compute the most likely path through our semantic type graph. 

The Viterbi algorithm consists of two phases -- forward and backward. In the forward phase, we move left to right, computing the log-probability of each transition at each step, as shown by the vectors below each position in the figure. When computing the transition from multiple nodes to a single node (such as the one from [T129, T116] to [T126], we compute for both paths and choose the maximum value. 

In the backward phase, we move from right to left, choosing the maximum probability node at each step. This is shown in the figure as boxed entries. We can then lookup the appropriate semantic type and return the most likely sequence of semantic types (shown in bold in the bottom of the figure). 

However, our objective is to return disambiguated concept linkages for entities. Given a disambiguated semantic type and multiple possibilities indicated by SciSpacy's linking process, we use the emission probabilities to choose the most likely concept to apply at the position. The result for our example is shown in the table below.

Entity-ID	Entity Span	Concept-ID	Concept Primary Name	Semantic Type Code	Semantic Type Name	Correct?
1	staining	C0487602	Staining method	T059	Laboratory Procedure	N/A*
2	antibodies	C0003241	Antibodies	T116	Amino Acid, Peptide, or Protein	Yes
3	cat	C0008169	Chloramphenicol O-Acetyltransferase	T116	Amino Acid, Peptide, or Protein	No
4	antigens	C0003320	Antigens	T129	Immunologic Factor	N/A*
5	binding	C1145667	Binding action	T052	Activity	Yes
6	antibodies	C0003241	Antibodies	T116	Amino Acid, Peptide, or Protein	Yes

(N/A: non-ambiguous mappings) 

I thought this might be an interesting technique to share, hence writing about it. In addition, in the spirit of reproducibility, I have also provided the following artifacts for your convenience.

1. Code: This github gist contains code that illustrates NER + NEL on an input sentence using SciSpacy and its UMLS integration, and then applies my adaptation of the Viterbi method (as described in this post) to disambiguate ambiguous entity linkages.
2. Data: I have also provided the transition and emission matrices, and their associated lookup tables, for convenience, as these can be time consuming to generate from scratch from the CORD-19 dataset.

As always, I appreciate your feedback. Please let me know if you find flaws with my approach, and/or you know of a better approach for entity disambiguation

Entity Co-occurrence graphs as Mind Maps

Some time ago, as part of a discussion I don't remember much about anymore, I was referred to this somewhat old (Jan/Feb 2018) set of articles about Deutsche Bank and its involvement in money laundering activities.

• DEUTSCHE BANK: A GLOBAL BANK FOR OLIGARCHS — AMERICAN & RUSSIAN, PART 1
• DEUTSCHE BANK: A GLOBAL BANK FOR OLIGARCHS — AMERICAN & RUSSIAN, PART 2
• DEUTSCHE BANK: A GLOBAL BANK FOR OLIGARCHS — AMERICAN & RUSSIAN, PART 3

Now I know as much about money laundering as the average person on the street, which is to say not much, so it was a fascinating and tedious read at the same time. Fascinating because of the scale of operations and the many big names involved, and tedious because there were so many players that I had a hard time keeping track of them as I read through the articles. In any case, I had just finished some work on my fork of the open source NERDS toolkit for training Named Entity Recognition models, and it occurred to me that identifying the entities in this set of articles and connecting them up into a graph might help to make better sense of it all. Sort of like how people draw mind-maps when trying to understand complex information. Except our process is going to be (mostly) automated, and our mind-map will have entities instead of concepts.

Skipping to the end, here is the entity graph I ended up building, it's a screenshot from the Neo4j web console. Red nodes represent persons, green nodes represent organizations, and yellow nodes represent geo-political entities. The edges are shown as directed, but of course co-occurrence relationships are bidirectional (or equivalently undirected).

The basic idea is to find Named Entities in the text using off the shelf Named Entity Recognizers (NERs), and connect a pair of entities if they co-occur in the same sentence. The transformation from unstructured text to entity graph is mostly automated, except for one step in the middle where we manually refine the entities and their synonyms. The graph data was ingested into a Neo4j graph database, and I used Cypher and Neo4j graph algorithms to generate insights from the graph. In this post I describe the steps to convert from unstructured article text to entity graph. The code is provided on GitHub, and so is the the data for this example, so you can use them to glean other interesting insights from this data, as well as rerun the pipeline to create entity graphs for your own text.

I structured the code as a sequence of Python scripts and Jupyter notebooks that are applied to the data. Each script or notebook reads the data files already available and writes new data files for the next stage. Scripts are numbered to indicate the sequence in which they should be run. I describe these steps below.

As mentioned earlier, the input is the text from the three articles listed above. I screen scraped the text into a local text file (select the article text and then copy the text, then paste it into a local text editor, and finally saved it into the file db-article.txt. The text is organized into paragraphs, with an empty line delimiting each paragraph. The first article also provided a set of acronyms and their expansions, which I captured similarly into the file db-acronyms.txt.

• 01-preprocess-data.py -- this script reads the paragraphs and converts it to a list of sentences. For each sentence, it checks to see if any token is an acronym, and if so, it replaces the token with the expansion. The script uses the SpaCy sentence segmentation model to segment the paragraph text into sentences, and the English tokenizer to tokenize sentences into tokens. Output of this step is a list of 585 sentences in the sentences.txt file.
• 02-find-entities.py -- this script uses the SpaCy pre-trained NER to find instances of Person (PER), Organization (ORG), GeoPolitical (GPE), Nationalities (NORP), and other types of entities. Output is written to the entities.tsv file, one entity per line.
• 03-cluster-entity-mentions.ipynb -- in this Jupyter notebook, we do simple rule-based entity disambiguation, so that similar entity spans found in the last step are clustered under the same entity -- for example, "Donald Trump", "Trump", and "Donald J. Trump", are all clustered under the same PER entity for "Donald J. Trump". The disambiguation finds similar spans of text (Jaccard token similarity) and considers those above a certain threshold to refer to the same entity. The most frequent entity types found are ORG, PERSON, GPE, DATE, and NORP. This step writes out each cluster as a key-value pair, with the key being the longest span in the cluster, and the value as a pipe-separated list of the other spans. Output from this stage are the files person_syns.csv, org_syns.csv, and gpe_syns.csv.
• 04-generate-entity-sets.py -- This is part of the manual step mentioned above. The *_syns.csv files contain clusters that are mostly correct, but because the clusters are based solely on lexical similarity, they still need some manual editing. For example, I found the "US Justice Department" and "US Treasury Department" in the same cluster, but "Treasury" in a different cluster. Similarly, "Donald J. Trump" and "Donald Trump, Jr." appeared in the same cluster. This script re-adjusts the clusters, removing duplicate synonyms for clusters, and assigning the longest span as the main entity name. It is designed to be run with arguments so you can version the *_syn.csv files. The repository contains my final manually updated files as gpe_syns-updated.csv, org_syns-updated.csv, and person_syns-updated.csv.
• 05-find-corefs.py -- As is typical in most writing, people and places are introduced in the article, and are henceforth referred to as "he/she/it", at least while the context is available. This script uses the SpaCy neuralcoref to resolve pronoun coreferences. We restrict the coreference context to the paragraph in which the pronoun occurs. Input is the original text file db-articles.txt and the output is a file of coreference mentions corefs.tsv. Note that we don't yet attempt to update the sentences in place like we did with the acronyms because the resulting sentences are too weird for the SpaCy sentence segmenter to segment accurately.
• 06-find-matches.py -- In this script, we use the *_syns.csv files to construct a Aho-Corasick Automaton object (from the PyAhoCorasick module), basically a Trie structure against which the sentences can be streamed. Once the Automaton is created, we stream the sentences against it, allowing it to identify spans of text that match entries in its dictionary. Because we want to match any pronouns as well, we first replace any coreferences found in the sentence with the appropriate entity, then run the updated sentence against the Automaton. Output at this stage is the matched_entities.tsv, a structured file of 998 entities containing the paragraph ID, sentence ID, entity ID, entity display name, entity span start and end positions.
• 07-create-graphs.py -- We use the keys of the Aho-Corasick Automaton dictionary that we created in the previous step to write out a CSV file of graph nodes, and the matched_entities.tsv to construct entity pairs within the same sentence to write out a CSV file of graph edges. The CSV files are in the format required by the neo4j-admin command, which is used to import the graph into a Neo4j 5.3 community edition database.
• 08-explore-graph.ipynb -- We have three kinds of nodes in the graph, PERson, ORGanization, and LOCation nodes. In this notebook, we compute PageRank on each type of node to find the top people, organizations, and locations we should look at. From there, we select a few top people and find their neighbors. One other feature we built was a search like functionality, where once two nodes are selected, we show a list of sentences where these two entities cooccur. And finally, compute the shortest path between a pair of nodes. The notebook shows the different queries, the associated Cypher queries (including calls to Neo4j Graph algorithms), as well as the outputs of these queries, its probably easier for you to click through and take a look yourself than for me to describe it.

There are obviously many other things that can be done with the graph, limited only by your imagination (and possibly by your domain expertise on the subject at hand). For me, the exercise was fun because I was able to use off the shelf NLP components (as opposed to having to train my own compoenent for my domain) to solve a problem I was facing. Using the power of NERs and graphs allows us to gain insights that would normally not be possible solely from the text.

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.