Training NER model with custom torch model

[carles11ds]

Hi, I am trying to train a custom NER model by integrating models with torch but I am having problems running the training even though I have followed the example in the documentation. Here is the config file, the model created with pytorch and registered in spacy and the command with the error when I try to train the model. I'm using spacy v3.5.0

CONFIG FILE

[paths]
train = null
dev = null
init_tok2vec = "es_core_news_lg"
vectors = null

[system]
gpu_allocator = null
seed = 0

[nlp]
lang = "es"
pipeline = ["tok2vec","ner"]
batch_size = 32
disabled = []
before_creation = null
after_creation = null
after_pipeline_creation = null
tokenizer = {"@tokenizers":"spacy.Tokenizer.v1"}

[components]

[components.ner]
factory = "ner"

[components.ner.model]
@architectures = "BiLSTM+CRF.v1"
nO = 3
nM = 200
nC = 3
dropout = 0.2
width = 96
embed_size = 200
hidden_width = 64
include_static_vectors = true

[components.tok2vec]
factory = "tok2vec"

[components.tok2vec.model]
@architectures = "spacy.Tok2Vec.v2"

[components.tok2vec.model.embed]
@architectures = "spacy.MultiHashEmbed.v2"
width = ${components.tok2vec.model.encode.width}
attrs = ["ORTH","SHAPE"]
rows = [5000,2500]
include_static_vectors = true

[components.tok2vec.model.encode]
@architectures = "spacy.MaxoutWindowEncoder.v2"
width = 256
depth = 8
window_size = 1
maxout_pieces = 3

[corpora]

[corpora.dev]
@readers = "spacy.Corpus.v1"
path = ${paths.dev}
gold_preproc = false
max_length = 0
limit = 0
augmenter = null

[corpora.train]
@readers = "spacy.Corpus.v1"
path = ${paths.train}
gold_preproc = false
max_length = 0
limit = 0
augmenter = null

[training]
accumulate_gradient = 1
dropout = 0.2
patience = 1600
max_epochs = 0
seed = ${system.seed}
gpu_allocator = ${system.gpu_allocator}
max_steps = 20000
eval_frequency = 200
frozen_components = []
annotating_components = []
before_to_disk = null
dev_corpus = "corpora.dev"
train_corpus = "corpora.train"
before_update = null

[training.batcher]
@batchers = "spacy.batch_by_words.v1"
discard_oversize = false
tolerance = 0.2
get_length = null

[training.batcher.size]
@schedules = "compounding.v1"
start = 100
stop = 1000
compound = 1.001
t = 0.0

[training.logger]
@loggers = "spacy.ConsoleLogger.v1"
progress_bar = false

[training.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = false
eps = 0.00000001
learn_rate = 0.001

[training.score_weights]
ents_f = 1.0
ents_p = 0.0
ents_r = 0.0
ents_per_type = null

[pretraining]

[initialize]
vectors = ${paths.vectors}
init_tok2vec = ${paths.init_tok2vec}
vocab_data = null
lookups = null
before_init = null
after_init = null

[initialize.components]

[initialize.tokenizer]

PYTORCH MODEL

from TorchCRF import CRF
from typing import List
from thinc.types import Floats2d
from thinc.api import Model, PyTorchWrapper, chain, with_array
import spacy
from spacy.tokens.doc import Doc
from spacy.ml import CharacterEmbed
from torch import nn

@spacy.registry.architectures("BiLSTM+CRF.v1")   
def create_torch_model(
    nO: int,
    width: int,
    hidden_width: int,
    embed_size: int,
    nM: int,
    nC: int,
    dropout: float,
    include_static_vectors: bool
) -> Model[List[Doc], List[Floats2d]]:
    char_embed = CharacterEmbed(width, embed_size, nM, nC, include_static_vectors)
    torch_model = nn.Sequential(
        nn.LSTM(embed_size, hidden_width, num_layers=2,
                              batch_first=True, dropout=dropout, bidirectional=True),
        nn.Linear(hidden_width * 2, nO),
        nn.Dropout2d(dropout),
        CRF(nO)
    )
    wrapped_pt_model = PyTorchWrapper(torch_model)
    model = chain(char_embed, with_array(wrapped_pt_model))
    return model

COMMAND

python -m spacy train config.cfg --code CRF_component.py --paths.dev test.spacy --paths.train train.spacy

ERROR

ℹ No output directory provided
ℹ Using CPU

=========================== Initializing pipeline ===========================
[2023-03-23 17:41:54,668] [INFO] Set up nlp object from config
[2023-03-23 17:41:54,674] [INFO] Pipeline: ['tok2vec', 'ner']
[2023-03-23 17:41:54,676] [INFO] Created vocabulary
[2023-03-23 17:41:54,677] [INFO] Finished initializing nlp object
Traceback (most recent call last):
  File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/spacy/__main__.py", line 4, in <module>
    setup_cli()
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/spacy/cli/_util.py", line 74, in setup_cli
    command(prog_name=COMMAND)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/typer/core.py", line 785, in main
    **extra,
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/spacy/cli/train.py", line 45, in train_cli
    train(config_path, output_path, use_gpu=use_gpu, overrides=overrides)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/spacy/cli/train.py", line 72, in train
    nlp = init_nlp(config, use_gpu=use_gpu)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/spacy/training/initialize.py", line 84, in init_nlp
    nlp.initialize(lambda: train_corpus(nlp), sgd=optimizer)
  File "/home/carlos/pycaret_env/lib/python3.7/site-packages/spacy/language.py", line 1308, in initialize
    proc.initialize(get_examples, nlp=self, **p_settings)
  File "spacy/pipeline/transition_parser.pyx", line 559, in spacy.pipeline.transition_parser.Parser.initialize
  File "spacy/pipeline/transition_parser.pyx", line 195, in spacy.pipeline.transition_parser.Parser._resize
KeyError: 'resize_output'

Any help would be appreciated

[kadarakos]

Hey carles11ds,

First of all let me start by saying that the ner component in spaCy is quite different than what you might expect. It is not like the usual sequence-tagging approach to Named Entity Recognition. For an introduction on the architecture please checkout: https://www.youtube.com/watch?v=sqDHBH9IjRU. It is a variation on the "transition-based" approach introduced in this paper: https://aclanthology.org/N16-1030/.

The concrete problem that you ran into is that the thinc model in the ner architecture is "resizeable". The ner model uses the "spacy.TransitionBasedParser.v2" architecture that you can find here, which internally uses this model to compute the best "action" for each step:

spaCy/spacy/ml/parser_model.pyx

Line 207 in d85df9d

class ParserStepModel(Model):

. This is where the resizing functionality is implemented: https://github.com/explosion/spaCy/blob/master/spacy/ml/models/parser.py#L101. This is useful when you would like to continue training your ner model with new classes.

If you are interested in implementing the BiLSTM-CRF I think an easier direction to pursue would be to implement a new TrainablePipe based on the tagger pipeline component: https://github.com/explosion/spaCy/blob/master/spacy/pipeline/tagger.pyx. This is a simple component that assigns a label for each token, much like the sequence-labeling ner approaches operate. Let's say your new component is called crf-tagger.

I would factor out the tok2vec part from you crf-tagger as well. I think it would be easier to follow the design of the tagger architecture: https://github.com/explosion/spaCy/blob/master/spacy/ml/models/tagger.py#L26. Here the tagger architecture itself does not define a new tok2vec only the classifier on top of the upstream tok2vec. In your config you can change the tok2vec separately from the crf-tagger. The architecture of the crf-tagger would just define the part downstream from the tok2vec meaning the linear-chain CRF sequence classifier. In the tok2vec layer you can use the "spacy.TorchBiLSTMEncoder.v1" that wraps the PyTorch LSTM you are using currently.

[carles11ds]

Hi kadarakos,

Thank you very much for your answer, I've been trying to include the CRF inside the tagger architecture but I'm having a lot of problems, as I don't know exactly what I have to change to get it to work. If you could help me with the main changes that need to be made in the tagger to define my crf_tagger component I would be very grateful.

[kadarakos]

Hey carles11ds,

It is a bit difficult for me to suggest complete path forward, but I can definitely highlight some hot spots.

First of all you have to decide on a tagging scheme which can be IOB or BILUO and spaCy has utilities for both of these: https://spacy.io/usage/linguistic-features#accessing-ner. Your crf tagger will take the entity annotation and convert them to token-level tags according to the tagging scheme of choice.

To me it seems like the main hurdle is being able to compute the right probabilities. Let me explain.
Currently the tagger uses the SequenceCategoricalCrossentropy loss from thinc as you can see here: https://github.com/explosion/spaCy/blob/master/spacy/pipeline/tagger.pyx#L261. As you can see this loss expects as guesses the type List[[Floats2d], where Floats2d is a 2d numpy or cupy array https://thinc.ai/docs/api-types#types. As explained in the docs for the CategoricalCrossentropy this Floats2d is expected contain the label probabilities.

So for each Doc the tagger model outputs a Floats2d with shape n_tokens by n_labels and this is used in the loss function. However, the TorchCRF model in the forward returns the FloatTensor of size batch_size instead: https://github.com/s14t284/TorchCRF/blob/master/TorchCRF/__init__.py#L52 containing the likelihood of the gold label sequence for each sequence. So for example if each element in the batch is a single Doc then you have a single scalar log-likelihood for each Doc, which is the log-likelihood the model assigned to the correct tag sequence.

I think the major change you would have to do is to reimplement forward to produce probabilities for each label for each token in each Doc. This part is a bit involved, but you can do it with the forward-backward algorithm that is partially implemented already in the library you are using. I would recommend this blogpost to learn more about the topic: https://lauraruis.github.io/2021/01/25/crfpt1.html.

Another potentially better alternative would be to figure out how to implement a thinc loss that can take the output of the forward function from TorchCRF and backpropagate properly. I'm tagging you @danieldk here if you'd have some opinions on what could be the most promising way of doing this.

We do not have time at the moment to prioritize the implementation of a linear-chain CRF in spaCy, but it is a feature that would be really nice to have for sure. We implement ideas like this in: https://github.com/explosion/spacy-experimental/