Create nlp pipeline utilizing custom ONNX Transformer plus NER pipe

[ricardoguedes0]

Hello everyone,

I have a fine-tuned transformer model for NER that will be going to production. As such, i need to reduce the model size and quantize it. For that, my go to is to utilize the ONNX format. The original model had a transformer, tagger and a ner in the pipeline. Not needing the tagger i removed it. Then, I extracted the transformer and convert it into the ONNX format successfully.
Right now my only problem is how to combine this new ONNX transformer model with the previous NER.
As of now i was able to create a Transformer component, inserted it in the pipeline and manage to run the code with the transformer embeddings being appended to the doc._.trf. I assuming that is enough for the NER TransformerListener to work however i still have no results. Here is my code so far.

def forward_onnx_model(self, docs, is_train=False):        
    token_doc = huggingface_tokenize(tokenizer = tokenizer, texts = [doc.text for doc in docs])
    wordpieces = WordpieceBatch(
        strings = token_doc["input_texts"],
        input_ids = token_doc["input_ids"],
        attention_mask = token_doc["attention_mask"].astype(np.float32),
        lengths = [len(doc) for doc in token_doc["input_texts"]],
        token_type_ids = token_doc["token_type_ids"],
    )
    # Calculate encodings
    encodings = {
        "input_ids": token_doc["input_ids"],
        "token_type_ids": token_doc["token_type_ids"],
        "attention_mask": token_doc["attention_mask"]
    }
    # output_names can be equal to None
    last_hidden_state, pooler_output = session.run(output_names=["output", "1711"], input_feed=encodings)
    model_output = ModelOutput([('last_hidden_state', last_hidden_state), ('pooler_output', pooler_output)])
    
    align = get_alignment(
        spans = [doc[:] for doc in docs],
        wordpieces= token_doc["input_texts"],
        special_tokens= list(set([token for tokens_doc in token_doc["input_texts"] for token in tokens_doc if token in tokenizer.special_tokens_map.values()]))
    )
    
    return [FullTransformerBatch(
            spans =  [doc[:] for doc in docs],
            wordpieces= wordpieces,
            model_output=model_output,
            align=align,
            cached_doc_data= [TransformerData(
                wordpieces=wordpieces,
                model_output=model_output,
                align=align
            )]
        )]


from spacy_transformers.annotation_setters import null_annotation_setter

@Language.factory("onnx_transformer")
def create_onnx_transformer(nlp: Language, name: str) -> Transformer:
    return  make_transformer(
    nlp= nlp,
    name= "onnx_transformer",
    model= model,
    max_batch_items= 4096,
    set_extra_annotations=null_annotation_setter
)

nlp.add_pipe("onnx_transformer", first=True)
print(nlp.pipeline)
[('onnx_transformer', <spacy_transformers.pipeline_component.Transformer object at 0x31d60b0a0>), ('ner', <spacy.pipeline.ner.EntityRecognizer object at 0x2bf40a490>)]

But with a simple example that should work. It fails to detect the Location "Portugal".
Example:

text = "Eu vivo em Portugal"
doc = nlp(text)
doc.to_json()
{'text': 'Eu vivo em Portugal',
 'ents': [],
 'tokens': [{'id': 0, 'start': 0, 'end': 2},
  {'id': 1, 'start': 3, 'end': 7},
  {'id': 2, 'start': 8, 'end': 10},
  {'id': 3, 'start': 11, 'end': 19}]}

My environment

spaCy version: 3.5.0
Platform: macOS-14.1.2-arm64-arm-64bit
Python version: 3.10.10