Is it possible to export and use a spaCy NER model without an vocab and inject tokens/vectors on the fly?

Viewed 263

INTRO

My question originates a bit far from the title but in essence sums up well what I'm stuck with at the moment.

I need to integrate a NER model of spacy as part of a complex and distributed NLP pipeline and what I'm doing to do so is:

  1. Train a new NER model based on en_core_web_lg model to recognize also my custom entities in NER task

  2. Save the model skipping the vocabulary to save disk space and memory usage

  3. Finally load the model to run some inference, using tokens and vectors that somebody pre-computed before in my pipeline, instead of compute it again using the model vocabulary (standard way).

The reason why I am saving the model without the vocab is because in my distributed pipeline one of the first things that is done is tokenize / vectorize the text so that the rest of the tasks have this input.

→ Before continuing I want to clarify that in the standard way (saving the vocab), I could train my custom NER, save / load and run an inference without major problems with a very good accuracy.

After that and reading the spaCy documentation, I found that it is possible to save my model without the vocabulary, and you can even build the Doc from a list of tokens and a custom vocabulary (in my case an empty vocab). Also, I was also able to set the document vectors using those that someone previously calculated for me in my pipeline.

However, when I save the model inside the ner/cfg file I see that there is a reference to the vectors on which the NER model was trained (en_core_web_lg.vectors):

{
   "disable":[
     "tagger",
     "parser"
   ],
   "beam_width":1,
   "beam_density":0.0,
   "beam_update_prob":1.0,
   "cnn_maxout_pieces":3,
   "nr_feature_tokens":6,
   "deprecation_fixes":{
      "vectors_name":"en_core_web_lg.vectors"
   },
   "nr_class":86,
   "hidden_depth":1,
   "token_vector_width":96,
   "hidden_width":64,
   "maxout_pieces":2,
   "pretrained_vectors":"en_core_web_lg.vectors",
   "bilstm_depth":0,
   "self_attn_depth":0,
   "conv_depth":4,
   "conv_window":1,
   "embed_size":2000
 }

This reference is the cause of an error when I try to load the model without having those vectors in memory (in other words without having that vocab loaded).

If I delete those references in the cfg file, the model loads correctly and I can run inferences using my vectors but the predictions obtained are very different from the ones I obtained with my first model (with the original vocab) and contain several errors.

QUESTION

This brings me to my original question: is it possible to save the NER model with an empty vocab and then run an inference with a SpaCy Doc built somehow from the tokens and vectors previously calculated in my pipeline?

Thanks a lot in advance!

BTW I'm using spaCy 2.3.7 and I put below some snippets of my code to clarify:

1. Training:

    nlp = spacy.load("en_core_web_lg", disable=['tagger', 'parser'])

    ner = nlp.get_pipe('ner')
    ner.add_label("FOO_ENTITY")
    ner.add_label("BAR_ENTITY")
    ner.add_label("COOL_ENTITY")

    # Start the training
    optimizer = nlp.begin_training()

    # Loop for EPOCHS iterations
    losses_hist = []
    for itn in range(30):
        # Shuffle the training data
        random.shuffle(Xy_train)
        losses = {}

        # Batch the examples and iterate over them
        for batch in spacy.util.minibatch(Xy_train, size=32):
            texts = [text for text, entities, _ in batch]
            golds = [{'entities': entities} for text, entities, _ in batch]

            # Update the model
            nlp.update(docs=texts, golds=golds, losses=losses)

        print(losses)
        losses_hist.append(losses)

2.a Run inference (standard):

    # I already have the text split in tokens
    doc = Doc(nlp.vocab, words=tokens)  # Create doc from tokens
    ner = nlp.get_pipe("ner")
    doc = ner(doc)  # Call NER step for doc

    for ent in doc.ents:
        print(f"value: {ent.text}, start: {ent.start_char}, end: {ent.end_char}, entity: {ent.label_}")

2.b Run inference with external vectors

    # I already have the text split in tokens and their vectors
    vectors = Vectors(data=embeddings, keys=tokens)
    nlp.vocab.vectors = vectors
    doc = Doc(nlp.vocab, words=tokens)
    
    ner = nlp.get_pipe("ner")
    doc = ner(doc)  # Call NER step for doc

    for ent in doc.ents:
        print(f"value: {ent.text}, start: {ent.start_char}, end: {ent.end_char}, entity: {ent.label_}")

3. Save / Load model:

# Save model
nlp.to_disk(str(dir_)

# Load model
nlp = spacy.load(str(dir_), exclude=['vocab']) 
2 Answers

In spacy v2 (not v3!) there are some hidden background steps that register the vectors globally under a particular name for use as features in the statistical models. (The idea behind this is that multiple models in the same process can potentially share the same vectors in RAM.)

To get a subset of vectors to work for a particular text, you need to register the new vectors under the right name:

  • use the same Vectors(name=) as in the model metadata when creating the vectors (will be something like en_core_web_lg.vectors)
  • run spacy._ml.link_vectors_to_models(vocab)

I'm pretty sure that this will start printing warnings and renaming the vectors internally based on the data shape if you do it repeatedly for different sets of vectors with the same name. I think you can ignore the warnings and it will work for that individual text, but it may break any other models loaded in the same script that are using that same vectors name/shape.

If you are doing this a lot in practice, you might want to write a custom version of link_vectors_to_models that iterates over the words in the vocab more efficiently for very small vector tables, or only modifies the words in the vocab that you know that you need. It really depends on the size of the vocab at the point where you're running link_vectors_to_models.

After several tests I discovered that the following way to reload the vectors for the vocab of my model in memory, works correctly for me. On the other hand this answer worked as inspiration but didn't work for spaCy 2.1.4.

The interesting thing is that:

  • I necessarily had to make a deep copy of the vocab, before updating it
  • The vectors must have exactly the same name as the vectors on which the NER model was trained (en_core_web_lg.vectors in my case)
  • It is necessary to update the vectors in memory as follows: thinc.extra.load_nlp.VECTORS
    import copy
    import spacy
    import thinc
    from spacy.vectors import Vectors
    from spacy.vocab import Vocab
    from thinc.v2v import Model
    
    vocab = copy.deepcopy(nlp.vocab)
    vectors = Vectors(data=embeddings, keys=tokens, name='en_core_web_lg.vectors')
    ops = Model.ops
    thinc.extra.load_nlp.VECTORS[(ops.device, 'en_core_web_lg.vectors')] = vectors.data  # Must re-load vectors
    vocab.vectors = vectors
    
    doc = Doc(vocab, words=tokens)  # Create doc from tokens an custom vocab
    ner = nlp.get_pipe("ner")
    doc = ner(doc)  # Call NER step for doc

I hope it can be of use to other users!

Related