Prepare training data and train custom NER using Spacy Python

In this post I will show you how to create final Spacy formatted training data to train custom NER using Spacy. And also show you how train custom NER by using this training data.

• Python version: 3.6.3
• Spacy version: 2.1.6
• en-core-web-sm (spacy small model) version: 2.1.0

###### Prepare Spacy formatted training data for custom NER #######

import json

# Read output json file from WebAnno (Annotation tool)
with open('input_json.json') as data_file:    
    data = json.load(data_file)

# Extract original sentences
sentences_list = data['_referenced_fss']['12']['sofaString'].split('rn')

# Extract entity start/ end positions and names
ent_loc = data['_views']['_InitialView']['NamedEntity']

# Extract Sentence start/ end positions
Sentence = data['_views']['_InitialView']['Sentence']

# Set first sentence starting position 0
Sentence[0]['begin'] = 0

# Prepare spacy formatted training data
TRAIN_DATA = []
ent_list = []
for sl in range(len(Sentence)):
    ent_list_sen = []
    for el in range(len(ent_loc)):
        if(ent_loc[el]['begin'] >= Sentence[sl]['begin'] and ent_loc[el]['end'] <= Sentence[sl]['end']):
            ## Need to subtract entity location with sentence begining as webanno generate data by treating document as a whole
            ent_list_sen.append([(ent_loc[el]['begin']-Sentence[sl]['begin']),(ent_loc[el]['end']-Sentence[sl]['begin']),ent_loc[el]['value']])
    ent_list.append(ent_list_sen)
    ## Create blank dictionary
    ent_dic = {}
    ## Fill value to the dictionary
    ent_dic['entities'] = ent_list[-1]
    ## Prepare final training data
    TRAIN_DATA.append([sentences_list[sl],ent_dic])
    
TRAIN_DATA

Now let’s try to train a new fresh NER model by using prepared custom NER data

import spacy
import random
from spacy.util import minibatch, compounding
from pathlib import Path

# Define output folder to save new model
model_dir = 'D:/Anindya/E/model'

# Train new NER model
def train_new_NER(model=None, output_dir=model_dir, n_iter=100):
    """Load the model, set up the pipeline and train the entity recognizer."""
    if model is not None:
        nlp = spacy.load(model)  # load existing spaCy model
        print("Loaded model '%s'" % model)
    else:
        nlp = spacy.blank("en")  # create blank Language class
        print("Created blank 'en' model")

    # create the built-in pipeline components and add them to the pipeline
    # nlp.create_pipe works for built-ins that are registered with spaCy
    if "ner" not in nlp.pipe_names:
        ner = nlp.create_pipe("ner")
        nlp.add_pipe(ner, last=True)
    # otherwise, get it so we can add labels
    else:
        ner = nlp.get_pipe("ner")

    # add labels
    for _, annotations in TRAIN_DATA:
        for ent in annotations.get("entities"):
            ner.add_label(ent[2])

    # get names of other pipes to disable them during training
    other_pipes = [pipe for pipe in nlp.pipe_names if pipe != "ner"]
    with nlp.disable_pipes(*other_pipes):  # only train NER
        # reset and initialize the weights randomly – but only if we're
        # training a new model
        if model is None:
            nlp.begin_training()
        for itn in range(n_iter):
            random.shuffle(TRAIN_DATA)
            losses = {}
            # batch up the examples using spaCy's minibatch
            batches = minibatch(TRAIN_DATA, size=compounding(4.0, 32.0, 1.001))
            for batch in batches:
                texts, annotations = zip(*batch)
                nlp.update(
                    texts,  # batch of texts
                    annotations,  # batch of annotations
                    drop=0.5,  # dropout - make it harder to memorise data
                    losses=losses,
                )
            print("Losses", losses)

    # test the trained model
    for text, _ in TRAIN_DATA:
        doc = nlp(text)
        print("Entities", [(ent.text, ent.label_) for ent in doc.ents])
        print("Tokens", [(t.text, t.ent_type_, t.ent_iob) for t in doc])

    # save model to output directory
    if output_dir is not None:
        output_dir = Path(output_dir)
        if not output_dir.exists():
            output_dir.mkdir()
        nlp.to_disk(output_dir)
        print("Saved model to", output_dir)

        # test the saved model
        print("Loading from", output_dir)
        nlp2 = spacy.load(output_dir)
        for text, _ in TRAIN_DATA:
            doc = nlp2(text)
            print("Entities", [(ent.text, ent.label_) for ent in doc.ents])
            print("Tokens", [(t.text, t.ent_type_, t.ent_iob) for t in doc])
            
# Finally train the model by calling above function
train_new_NER()

After running above code you should find that some files are created in the specified folder.

Test new trained NER model in Spacy

Now it’s time to test our fresh trained NER model to see whether it is working properly or not.

# Use new trained saved model
print("Loading trained model from:", model_dir)
nlp2 = spacy.load(model_dir)
doc2 = nlp2('Who is Shaka Khan?')

for token in doc2:
    print(token, token.ent_type_)

updated_model_dir = 'D:/Anindya/E/updated_model'

## Update existing spacy model and store into a folder
def update_model(model='en_core_web_sm', output_dir=updated_model_dir, n_iter=100):
    """Load the model, set up the pipeline and train the entity recognizer."""
    if model is not None:
        nlp = spacy.load(model)  # load existing spaCy model
        print("Loaded model '%s'" % model)
    else:
        nlp = spacy.blank("en")  # create blank Language class
        print("Created blank 'en' model")

    # create the built-in pipeline components and add them to the pipeline
    # nlp.create_pipe works for built-ins that are registered with spaCy
    if "ner" not in nlp.pipe_names:
        ner = nlp.create_pipe("ner")
        nlp.add_pipe(ner, last=True)
    # otherwise, get it so we can add labels
    else:
        ner = nlp.get_pipe("ner")

    # add labels
    for _, annotations in TRAIN_DATA:
        for ent in annotations.get("entities"):
            ner.add_label(ent[2])

    # get names of other pipes to disable them during training
    other_pipes = [pipe for pipe in nlp.pipe_names if pipe != "ner"]
    with nlp.disable_pipes(*other_pipes):  # only train NER
        # reset and initialize the weights randomly – but only if we're
        # training a new model
        if model is None:
            nlp.begin_training()
        for itn in range(n_iter):
            random.shuffle(TRAIN_DATA)
            losses = {}
            # batch up the examples using spaCy's minibatch
            batches = minibatch(TRAIN_DATA, size=compounding(4.0, 32.0, 1.001))
            for batch in batches:
                texts, annotations = zip(*batch)
                nlp.update(
                    texts,  # batch of texts
                    annotations,  # batch of annotations
                    drop=0.5,  # dropout - make it harder to memorise data
                    losses=losses,
                )
            print("Losses", losses)

    # test the trained model
    for text, _ in TRAIN_DATA:
        doc = nlp(text)
        print("Entities", [(ent.text, ent.label_) for ent in doc.ents])
        print("Tokens", [(t.text, t.ent_type_, t.ent_iob) for t in doc])

    # save model to output directory
    if output_dir is not None:
        output_dir = Path(output_dir)
        if not output_dir.exists():
            output_dir.mkdir()
        nlp.to_disk(output_dir)
        print("Saved model to", output_dir)

# Finally train the model by calling above function        
update_model()

After running above code you should find that some files are created in the specified folder.

 

Test updated NER model in Spacy

Now it’s time to test our updated NER model to see whether it is working properly or not.

# Use updated saved model
print("Loading updated model from:", updated_model_dir)
nlp2 = spacy.load(model_dir)
doc2 = nlp2('Who is Shaka Khan?')

for token in doc2:
    print(token, token.ent_type_)

• How to create Spacy formatted training data for custom NER
• Train Custom NER model using Spacy in python
• Update existing Spacy NER model