The best way to combine en_coreference_web_trf with en_core_web_trf in a single pipeline

[sztal]

Hi,
first of all, thanks for the awesome work with spacy! Now, I am working on an application which depends on both super accurate and meaningful linguistic annotations (.pos and .dep) provided by en_core_web_trf model (which works really well for my usecase) as well as coreference resolution for which I want to use the new en_coreference_web_trf model.

However, it seems both of them use different transformer pipeline components and when I try to run coref component based on transformer from en_core_web_trf the results are significantly worse. Similarily, when I use transformer from en_coreference_web_trf the linguistic annotation produced by en_core_web_trf are really meaningless and useless and do not make any sense.

Is there any way I can have the best of both worlds? Below is an example of a the problem.

Case I: using transformer from en_core_web_trf

import spacy
from spacy.lang.en import English

text = "Philip plays the bass because he loves it."

# NLP = spacy.load("en_core_web_trf")
COREF = spacy.load("en_coreference_web_trf")
COREF.disable_pipes("span_resolver", "span_cleaner")
CORE = spacy.load("en_core_web_trf")

NLP = English()
for pipe in CORE.pipe_names:
    NLP.add_pipe(pipe, source=CORE)
for pipe in COREF.pipe_names:
    if pipe in NLP.pipe_names:
        continue
    NLP.add_pipe(pipe, source=COREF)

doc = NLP(text)
doc.spans
# OUTPUT: {'coref_head_clusters_1': [bass, it]}
# 'Philip' and 'he' are not clustered

Case II: using transformer from en_coreference_web_trf

import spacy
from spacy.lang.en import English

text = "Philip plays the bass because he loves it."

# NLP = spacy.load("en_core_web_trf")
COREF = spacy.load("en_coreference_web_trf")
COREF.disable_pipes("span_resolver", "span_cleaner")
CORE = spacy.load("en_core_web_trf")

NLP = English()
for pipe in COREF.pipe_names:
    NLP.add_pipe(pipe, source=COREF)
for pipe in CORE.pipe_names:
    if pipe in NLP.pipe_names:
        continue
    NLP.add_pipe(pipe, source=CORE)

doc = NLP(text)
doc.spans
# OUTPUT: {'coref_head_clusters_1': [Philip, he], 'coref_head_clusters_2': [bass, it]}
# Now clusters are okay
# But linguistic annotations do not make sense
[ (t.i, t.text, t.pos_, t.dep_) for t in doc ]
# OUTPUT:
# [(0, 'Philip', 'PROPN', 'ROOT'),
#  (1, 'plays', 'AUX', 'aux'),
#  (2, 'the', 'PRON', 'prep'),
#  (3, 'bass', 'AUX', 'aux'),
#  (4, 'because', 'PROPN', 'nsubj'),
#  (5, 'he', 'PROPN', 'nsubj'),
#  (6, 'loves', 'VERB', 'advcl'),
#  (7, 'it', 'PRON', 'prep'),
#  (8, '.', 'PUNCT', 'advcl')]
# For instance 'because' is marked as proper noun and nominal subject

[svlandeg]

first of all, thanks for the awesome work with spacy!

Thanks! 🤗

Sorry that this has been inconvenient. We appreciate reports of use-cases like yours to help us establish what functionality is needed when we pull the coref component out of experimental. When we do, we definitely want to make sure this new component is more easily usable with other trained components, e.g. by ensuring they're trained on the same transformer.

For now, if you want to combine these models, you can use the function replace_listeners which moves a shared transformer layer into the component. In general, that would look something like this:

    import spacy

    # move the 'transformer' from the coref pipeline to each component, then remove it
    nlp = spacy.load("en_coreference_web_trf")
    nlp.replace_listeners("transformer", "coref", ["model.tok2vec"])
    nlp.replace_listeners("transformer", "span_resolver", ["model.tok2vec"])
    nlp.remove_pipe("transformer")

    # copy all core components (including the transformer) as is
    CORE = spacy.load("en_core_web_trf")
    for pipe in CORE.pipe_names:
        nlp.add_pipe(pipe, source=CORE)

    doc = nlp("Philip plays the bass because he loves it.")
    print(doc.spans)
    print([(t.i, t.text, t.pos_, t.dep_) for t in doc])

Which should give you sensible results for all components.

[sztal]

Thanks! I will soon try to implement this solution and see how it works. In the meantime I did a workaround based on parsing each document two times (once with core model and once with coreference model) and then aligning tokenizations. Below is a sketch of the main methods of the corresponding pipeline component, in case someone is interested.

def __call__(self, doc: Doc) -> Doc:
        cdoc = self.model(doc.text)  # this stores `en_coreference_web_trf`
        s1 = list(map(str, doc))
        s2 = list(map(str, cdoc))
        align = Alignment.from_strings(s1, s2)  # from spacy.training import Alignment
        for spans in cdoc.spans.values():
            cluster = [ int(align.y2x.data[t.i]) for s in spans for t in s ]
            self.set_corefs(doc, cluster)
        return doc

    def set_corefs(self, doc: Doc, cluster: Sequence[int]) -> None:
        """Set proper coreferences from pronoun tokens to closest
        non-pronoun neighbors within the ``cluster``.
        """
        alias = "some_alias_for_extensions_attributes"
        proper = []
        pronouns = []
        for i in cluster:
            if doc[i].pos_ == "PRON":
                pronouns.append(i)
            else:
                proper.append(i)
        for i in pronouns:
            closest = None
            for j in proper:
                if closest is None or abs(i-j) < closest:
                    closest = j
            if closest is not None:
                refs = set()
                refs.add(closest)
                for conj in doc[closest].conjuncts:
                    refs.add(conj.i)
                setattr(doc[i]._, f"{alias}_refs", tuple(sorted(refs)))

[sztal]

And just one more question. I assume that your solution will be much faster as it will require only a single document parsing, right?

[adrianeboyd]

Combining these models into one pipeline in the most efficient way is tricky due to the some of the internal transformer listener settings, so the solution above is a bit larger and slower than it ought to be because there are effectively three transformer models.

What I would recommend given the current state of affairs and so that you don't have to align tokens across docs: load both pipelines separately, but with the same vocab so you can run them on the same docs (only two transformer models rather than three):

nlp = spacy.load("en_core_web_trf")
nlp_coref = spacy.load("en_coreference_web_trf", vocab=nlp.vocab)

doc = nlp(text)

# note that here you're passing in the doc rather than the text
doc = nlp_coref(doc)

We will work on making this easier in the future! As a first step, the model en_coreference_web_trf could use a custom transformer component name, and then you could simply add all the components to the same pipeline.

[sztal]

Hi, thanks a lot! I will try this and let you know if it works as expected in our case.

[Jiya126]

I'm printing the POS and dependency relation using the command you used above i.e.

print([ (t.i, t.text, t.pos_, t.dep_) for t in doc ])

but i'm getting output to be

[(0, 'The', '', ''), (1, 'cats', '', ''), (2, 'were', '', ''), (3, 'startled', '', ''), (4, 'by', '', ''), (5, 'the', '', ''), (6, 'dog', '', ''), (7, 'as', '', ''), (8, 'it', '', ''), (9, 'growled', '', ''), (10, 'at', '', ''), (11, 'them', '', ''), (12, '.', '', '')]

Here's the whole code for reference

import spacy
nlp = spacy.load("en_coreference_web_trf")
doc = nlp("The cats were startled by the dog as it growled at them.")
print(doc.spans)

print([ (t.i, t.text, t.pos_, t.dep_) for t in doc ])

I'm using only "en_coreference_web_trf" here. Also, it is printing the coref_clusters from doc.spans but not pos and dep.

[Jiya126]

@svlandeg @sztal

[rmitsch]

Are you referring to the command/snippet in this comment? Also - please don't ping maintainers directly. We see your posts and will respond to them without being pinged.

[Jiya126]

Sorry for that
Yes, that is the command i'm referring to

[rmitsch]

The snippet you have under "Here's the whole code for reference" is different from the snippet in the comment I linked though. Note that you're using a differently configured model. What output do you get for running the original, linked code snippet?