Trainable pipeline with pre-added layout spans

[thatbudakguy]

Maybe this isn't the right way to do this, but I'm attempting to run the layout before the rest of my pipeline, so that I can use the spans identified by it and pass them to my components.

Specifically, the use case is taking the identified layout spans (Doc.spans["layout"]) and using them as the suggested spans for the SpanCategorizer directly (via spacy.preset_spans_suggester.v1).

What I tried to do was define a callback like:

# Callback that allows adding layout to the pipeline during training
@spacy.registry.callbacks("add_layout")
def create_add_layout_callback():
    def add_layout(nlp):
        from spacy_layout import spaCyLayout
        return spaCyLayout(nlp)
    return add_layout

And use it in my config, e.g.:

[nlp]
lang = "en"
pipeline = ["tok2vec","spancat"]
batch_size = 1000
disabled = []
before_creation = null
after_creation = {"@callbacks": "add_layout"}
after_pipeline_creation = null
tokenizer = {"@tokenizers":"spacy.Tokenizer.v1"}
vectors = {"@vectors":"spacy.Vectors.v1"}

And then invoke spacy train --code my_code.py to pick up the function.

What I get is:

ValueError: [E942] Executing after_creation callback failed. Expected the function to return an initialized nlp object but got: <class 'spacy_layout.layout.spaCyLayout'>. Maybe you forgot to return the modified object in your function?

[thatbudakguy]

...I guess this is a bad idea, since spaCyLayout actually expects a path to a file rather than text.

What I'm interested in ultimately is having layout information already annotated on the training docs. I generated and exported annotations using prodigy's pdf.spans.manual recipe, but it seems like the data I exported with data-to-spacy does not ultimately include the Doc.spans["layout"] information, because when I try to train, I see KeyError: 'layout'.

[thatbudakguy]

It seems like the exported data actually treats each span as a doc, so this strategy won't work.