spaCy: How to update doc.ents when using doc.retokenize()

Viewed 99

I am trying to update a pre-trained model with tokens using the retokenizer. I created a pipeline in order to do this. In this pipeline, I also set "ENT_TYPE" when merging the tokens.

@Language.factory("re_tokenize")
def re_tokenize(nlp, name):
    return ReTokenize(nlp.vocab)

class ReTokenize:
    pattern = ""
    
    def __init__(self, vocab):
        self.pattern = r"[a-zA-Z0-9]+\[{0,1}[a-zA-Z0-9_]+\]{0,1}\[{0,1}[a-zA-Z0-9_]+\]{0,1}\[{0,1}[a-zA-Z0-9_]+\]{0,1}#{0,1}"
    
    def __call__(self, doc):
        spans = []

        for match in re.finditer(self.pattern, doc.text):
            start, end = match.span()
            span = doc.char_span(start, end)

            if span is not None:
                spans.append(span)

        with doc.retokenize() as retokenizer:
            for span in spans:
                retokenizer.merge(span, attrs={"ENT_TYPE": "VAR"})

        return doc

Using this pipeline, I can tokenize the words correctly. Also, the data in ent_type_ seems to be updated.

BEFORE:

# Set model
nlp = spacy.load("ja_ginza")

text = "aaa_bbbとaaa_CCCの2バイトマップ"
text = mojimoji.zen_to_han(text).lower()

doc = nlp(text)

print([token.text for token in doc])
print([token.ent_type_ for token in doc])

['aaa', '', 'bbb', 'と', 'aaa', '', 'ccc', 'の', '2', 'バイト', 'マップ']

['Product_Other', 'Product_Other', 'Product_Other', '', 'Product_Other', 'Product_Other', 'Product_Other', '', 'N_Product', 'N_Product', 'N_Product']

AFTER:

nlp.add_pipe("re_tokenize", before="parser")
doc = nlp(text)

print([token.text for token in doc])
print([token.ent_type_ for token in doc])

['aaa_bbb', 'と', 'aaa_ccc', 'の', '2', 'バイト', 'マップ']

['VAR', '', 'VAR', '', 'N_Product', 'N_Product', 'N_Product']

However, it seems that doc.ents is not being updated:

print([ent.label_ for ent in doc.ents])

['N_Product']

How do I also update doc.ents?

1 Answers

To add a single new entity to a doc without modifying any other entity annotation, use doc.set_ents():

span = doc.char_span(start, end, label="VAR")
doc.set_ents(entities=[span], default="unmodified")

More docs: https://spacy.io/api/doc#set_ents

Related