Integrating custom NER models into spaCy pipeline

[admatis]

Hi all, we are trying to wrap our Hebrew pytorch custom transformer-based model as spacy model. We currently have .pt file for the finetuned transformer (that can be loaded using 🤗) and two .pt files each for token-based classification of BIO tags (each responsible to a different set of entity types). We wondered what is the quickest/cleanest way to wrap our setup as a spacy model to enable easy use.

We figured there should be a pipe to load the transformer as an embedding model and add two NER model pipes (for each BIO classifier) as listeners to the embeddings and let spacy handle the conversion from BIO over BPE tokens to entity spans.

We managed to load the transformer as a pipe (using a transformer-component config file) but got stuck adding the NER classification heads.

Any help would be greatly appreciated.

---

Our current code for loading the Hebrew finetuned transformer (tokenizer - onlplab/alephbert-base)

from thinc.config import Config
from spacy.lang.xx import MultiLanguage

DEFAULT_TRANSFORMER_CONFIG_STR = """
[transformer]
max_batch_items = 4096

[transformer.set_extra_annotations]
@annotation_setters = "spacy-transformers.null_annotation_setter.v1"

[transformer.model]
@architectures = "spacy-transformers.TransformerModel.v3"
name = "./finetuned_alephbert"
tokenizer_config = {"use_fast": true}
transformer_config = {}
mixed_precision = false
grad_scaler_config = {}

[transformer.model.get_spans]
@span_getters = "spacy-transformers.strided_spans.v1"
window = 128
stride = 96
"""
DEFAULT_TRANSFORMER_CONFIG = Config().from_str(DEFAULT_TRANSFORMER_CONFIG_STR)
nlp = MultiLanguage()

transformer = nlp.add_pipe("transformer", config=DEFAULT_TRANSFORMER_CONFIG["transformer"])
transformer.model.initialize()
text = "שאול עולה לכיתה א'"
print(nlp.components)
doc = nlp(text)
tokvecs = doc._.trf_data.tensors[-1]

[polm]

Sorry for the delayed reply on this! Unfortunately our standard reply here is that spacy-transformers simply doesn't support task-specific heads. You can use the Transformer as a source of features and train spaCy native NER layers.

It may be possible to implement a workaround wrapping the model in Thinc, though I think it'd be pretty involved and I'm not sure anyone has done that before.

Another thing you can do is keep the Transformer model separate from spaCy, and just use the annotations from it to create Docs manually. That's pretty inefficient and loses a lot of the benefits of spaCy, so we don't generally recommend it, but it can still be worthwhile if you have a lot of postprocessing that would benefit from spaCy's architecture.