Pass additional arguments through Language.__call__/pipe to Language.make_doc

[kinghuang]

I've created a Language subclass to extend the make_doc method with some additional logic. make_doc is normally invoked by Language.__call__ and Language.pipe to create a Doc from text.

To accomplish my goals, I need to pass additional information to make_doc. However, there isn't a way to do this via __call__ or pipe. Language.pipe does have the notion of a context, but this is simply yielded to the caller.

I propose that __call__ and make_doc be modified to accept additional args and/or keyword args, that __call__ pass the additional args/kwargs to make_doc, and that pipe pass the context per text to make_doc.

Example Code

The current signature for Language.__call__ and the invocation of make_doc is:

def __call__(
    self,
    text: str,
    *,
    disable: Iterable[str] = SimpleFrozenList(),
    component_cfg: Optional[Dict[str, Dict[str, Any]]] = None,
) -> Doc:

doc = self.make_doc(text)

I've extended this and make_doc in my custom Language subclass to accept and pass additional arguments and keyword arguments.

def __call__(
    self,
    text: str,
    *args,
    disable: Iterable[str] = SimpleFrozenList(),
    component_cfg: Optional[Dict[str, Dict[str, Any]]] = None,
    **kwargs,
) -> Doc:
    doc = self.make_doc(text, *args, **kwargs)
    
def make_doc(self, text: str, *args, **kwargs) -> Doc:
    pass

For Language.pipe, the context associated with each text could be passed to make_doc. I haven't prototyped this because it'd involve a bigger change and there's extra considerations required for the multiprocessing pipeline.

Use Case

One of my use cases for this is custom affix handling. I have some text preprocessing that currently strips away affixes before the text goes into a spaCy language pipeline. The affixes are different for each piece of text, so they can't be generalized into an affix pattern for the tokenizer. Here's some made up examples.

affix1 = "ABCDE"
text1 = "The quick brown foxABCDE"

affix2 = "12345"
text2 = "123451 Infinite Loop"

Because the affixes are unique (and not nearly as simple as this in the actual data!), I want to pass that to make_docs, and store the data on a Doc extension. That way, I can have a custom component read the affix and retokenize the Doc and perform other operations on the affixes.

Doc.set_extension("affix", default=None)

class CustomEnglishLanguage(English):

    def make_doc(self, text: str, affix: str = None, *, **) -> Doc:
        doc = super().make_doc(text)

        if affix:
            doc._.set("affix", affix)

        return doc

doc1 = nlp("The quick brown foxABCDE", affix="ABCDE")
doc2 = nlp("123451 Infinite Loop", affix="12345")

Rationale

The main reason why additional arguments to make_doc is needed is that the custom pipeline components need additional information about the text, such as the affix in the above example, on a per text basis. Setting the information on the Doc returned by __call__/pipe is too late, because the Doc has already passed through the language pipeline. The make_doc method is the ideal place to set additional information because it constructs the Doc that is passed through the language pipeline, hence additional information can be set and the Doc otherwise customized before pipeline components run.

At the moment, I have this working in my project. But, it's not ideal because I have to copy the entire implementation of Language.__call__ just to change the call to make_doc.

[adrianeboyd]

There's a definite need for this kind of functionality. We're currently planning to do this slightly differently by letting you just pass Doc objects to the pipeline, and the pipeline will skip make_doc if it gets a Doc instead of a text. Your code would have the steps:

• create a Doc (with nlp.make_doc in many cases, but you can use any external tokenizer / data source instead)
• add metadata, partial annotation, etc.
• call nlp(doc) or nlp.pipe(docs)

There is also always the option to do this with a custom tokenizer so you can have a single pipeline call instead of multiple steps, but it's much more of a hassle if you want to extend the built-in Tokenizer class due to the use of cython. If you have some other python-only tokenizer, it could be pretty straightforward, but most languages do use the Tokenizer.

The draft PR: #9069

[kinghuang]

Ha, that's great timing! Your approach makes sense to me. Subclassing Language and Tokenizer (which I've also done for other reasons) is not convenient. Being able to pass Doc objects to __call__ and pipe would be really useful.