How Do I Serialize spaCy Custom Span Extensions as JSON? - json

Question

I am using spaCy 2.1.6 to define a custom extension on a span.
>>> from spacy import load
>>> nlp = load("en_core_web_lg")
>>> from spacy.tokens import Span
>>> Span.set_extension('my_label', default=None)
>>> d = nlp("The fox jumped.")
>>> d[0:2]._.my_label = "ANIMAL"
>>> d[0:2]._.my_label
'ANIMAL'
The custom span extension does not appear when I serialize the document to JSON.
>>> d.to_json()
{'text': 'The fox jumped.',
'ents': [],
'sents': [{'start': 0, 'end': 15}],
'tokens': [{'id': 0,
'start': 0,
'end': 3,
'pos': 'DET',
'tag': 'DT',
'dep': 'det',
'head': 1},
{'id': 1,
'start': 4,
'end': 7,
'pos': 'NOUN',
'tag': 'NN',
'dep': 'nsubj',
'head': 2},
{'id': 2,
'start': 8,
'end': 14,
'pos': 'VERB',
'tag': 'VBD',
'dep': 'ROOT',
'head': 2},
{'id': 3,
'start': 14,
'end': 15,
'pos': 'PUNCT',
'tag': '.',
'dep': 'punct',
'head': 2}]}
(I'm specifically interested in custom annotation of Spans, but the same appears to be true of the JSON serialization of Doc object.)
Pickling and unpickling the document does preserve the custom extension.
How do I get the custom span extensions into the JSON serialization, or is that not supported?

Answer 1

Use this function and add your custom extensions any way you want:
def doc2json(doc: spacy.tokens.Doc, model: str):
json_doc = {
"text": doc.text,
"text_with_ws": doc.text_with_ws,
"cats": doc.cats,
"is_tagged": doc.is_tagged,
"is_parsed": doc.is_parsed,
"is_nered": doc.is_nered,
"is_sentenced": doc.is_sentenced,
}
ents = [
{"start": ent.start, "end": ent.end, "label": ent.label_} for ent in doc.ents
]
if doc.is_sentenced:
sents = [{"start": sent.start, "end": sent.end} for sent in doc.sents]
else:
sents = []
if doc.is_tagged and doc.is_parsed:
noun_chunks = [
{"start": chunk.start, "end": chunk.end} for chunk in doc.noun_chunks
]
else:
noun_chunks = []
tokens = [
{
"text": token.text,
"text_with_ws": token.text_with_ws,
"whitespace": token.whitespace_,
"orth": token.orth,
"i": token.i,
"ent_type": token.ent_type_,
"ent_iob": token.ent_iob_,
"lemma": token.lemma_,
"norm": token.norm_,
"lower": token.lower_,
"shape": token.shape_,
"prefix": token.prefix_,
"suffix": token.suffix_,
"pos": token.pos_,
"tag": token.tag_,
"dep": token.dep_,
"is_alpha": token.is_alpha,
"is_ascii": token.is_ascii,
"is_digit": token.is_digit,
"is_lower": token.is_lower,
"is_upper": token.is_upper,
"is_title": token.is_title,
"is_punct": token.is_punct,
"is_left_punct": token.is_left_punct,
"is_right_punct": token.is_right_punct,
"is_space": token.is_space,
"is_bracket": token.is_bracket,
"is_currency": token.is_currency,
"like_url": token.like_url,
"like_num": token.like_num,
"like_email": token.like_email,
"is_oov": token.is_oov,
"is_stop": token.is_stop,
"is_sent_start": token.is_sent_start,
"head": token.head.i,
}
for token in doc
]
return {
"model": model,
"doc": json_doc,
"ents": ents,
"sents": sents,
"noun_chunks": noun_chunks,
"tokens": tokens,
}

Answer 2

Since I ran into the same issue and the only other answer didnt really help my I thought I mide as well give other persons looking into this some hints.
Since Spacy 2.1 Spacy removed print_tree and added the to_json. to_json does not return custom extensions as "this method will output the same format as the JSON training data expected by spacy train" (https://spacy.io/usage/v2-1).
If you want to output your custom extension you need to write your own to_json function.
To do this I recommend extending the to_json() given by spacy.

Answer 3

Not really a fan of the other two answers here since they seem a bit overkill (extending the Doc object by #Chooklii or the custom but flaky doc2json method solution by #Laksh) so I'll just drop here what I did for one of my projects here and maybe that is useful to someone.
doc = <YOUR_DOC_OBJECT>
extra_fields = [field for field in dir(doc._) if field not in ('get', 'set', 'has')]
doc_json = doc.to_json()
doc_json.update({field: doc._.get(field) for field in extra_fields})
The doc_json should now have all the fields that you set via the Extensions interface provided by spaCy along with the fields set by other spaCy pipelines.