Add ability to serialize pipeline with spacy

README.md

@@ -74,6 +74,48 @@ The following extension properties are available:
 
 These methods will raise an exception when called on a span that is not a constituent in the parse tree. Such errors can be avoided by traversing the parse tree starting at either sentence level (by iterating over `doc.sents`) or with an individual `Token` object.
 
+## Serialize pipeline with SpaCy
+
+### to_disk
+
+```python
+from spacy.tokens import DocBin
+
+doc = nlp("the little lion is sleeping")
+
+# save doc to file
+doc_bin = DocBin(store_user_data=True)
+doc_bin.add(doc)
+doc_bin.to_disk("./serialized.doc")
+
+# load file from disk
+doc_bin = DocBin().from_disk("./serialized.doc")
+restored_doc = list(doc_bin.get_docs(nlp.vocab))[0]
+
+list(restored_doc.sents)[0]._.parse_string
+# > '(SENT (NP (X the) (X little) (NOUN lion)) (VN (X is)) (AP (X sleeping)))'
+```
+
+### to_bytes
+```python
+from spacy.tokens import DocBin
+
+doc = nlp("the little lion is sleeping")
+
+# save doc to file
+doc_bin = DocBin(store_user_data=True)
+doc_bin.add(doc)
+_bytes = doc_bin.to_bytes()
+
+# load file from bytes
+doc_bin = DocBin().from_bytes(_bytes)
+restored_doc = list(doc_bin.get_docs(nlp.vocab))[0]
+
+list(restored_doc.sents)[0]._.parse_string
+# > '(SENT (NP (X the) (X little) (NOUN lion)) (VN (X is)) (AP (X sleeping)))'
+```
+
+
 ### Usage with NLTK
 
 There is also an NLTK interface, which is designed for use with pre-tokenized datasets and treebanks, or when integrating the parser into an NLP pipeline that already performs (at minimum) tokenization and sentence splitting. For parsing starting with raw text, it is **strongly encouraged** that you use spaCy and `benepar.BeneparComponent` instead.

src/benepar/integrations/spacy_extensions.py

@@ -13,6 +13,14 @@ def __init__(self, starts, ends, labels, loc_to_constituent, label_vocab):
         self.loc_to_constituent = loc_to_constituent
         self.label_vocab = label_vocab
 
+    def serialize(self):
+        return {
+            "starts": self.starts,
+            "ends": self.ends,
+            "labels": self.labels,
+            "loc_to_constituent": self.loc_to_constituent,
+            "label_vocab": self.label_vocab
+        }
 
 def get_constituent(span):
     constituent_data = span.doc._._constituent_data
@@ -22,15 +30,15 @@ def get_constituent(span):
             " Consider adding a BeneparComponent to the pipeline."
         )
 
-    search_start = constituent_data.loc_to_constituent[span.start]
-    if span.start + 1 < len(constituent_data.loc_to_constituent):
-        search_end = constituent_data.loc_to_constituent[span.start + 1]
+    search_start = constituent_data["loc_to_constituent"][span.start]
+    if span.start + 1 < len(constituent_data["loc_to_constituent"]):
+        search_end = constituent_data["loc_to_constituent"][span.start + 1]
     else:
-        search_end = len(constituent_data.ends)
+        search_end = len(constituent_data["ends"])
     found_position = None
     for position in range(search_start, search_end):
-        if constituent_data.ends[position] <= span.end:
-            if constituent_data.ends[position] == span.end:
+        if constituent_data["ends"][position] <= span.end:
+            if constituent_data["ends"][position] == span.end:
                 found_position = position
             break
 
@@ -41,13 +49,15 @@ def get_constituent(span):
 
 def get_labels(span):
     constituent_data, position = get_constituent(span)
-    label_num = constituent_data.labels[position]
-    return constituent_data.label_vocab[label_num]
+    label_num = constituent_data["labels"][position]
+    return constituent_data["label_vocab"][label_num]
 
+def get_token_labels(token):
+    return get_labels(token.doc[token.i : token.i + 1])
 
 def parse_string(span):
     constituent_data, position = get_constituent(span)
-    label_vocab = constituent_data.label_vocab
+    label_vocab = constituent_data["label_vocab"]
     doc = span.doc
 
     idx = position - 1
@@ -56,9 +66,9 @@ def make_str():
         nonlocal idx
         idx += 1
         i, j, label_idx = (
-            constituent_data.starts[idx],
-            constituent_data.ends[idx],
-            constituent_data.labels[idx],
+            constituent_data["starts"][idx],
+            constituent_data["ends"][idx],
+            constituent_data["labels"][idx],
         )
         label = label_vocab[label_idx]
         if (i + 1) >= j:
@@ -77,9 +87,9 @@ def make_str():
         else:
             children = []
             while (
-                (idx + 1) < len(constituent_data.starts)
-                and i <= constituent_data.starts[idx + 1]
-                and constituent_data.ends[idx + 1] <= j
+                (idx + 1) < len(constituent_data["starts"])
+                and i <= constituent_data["starts"][idx + 1]
+                and constituent_data["ends"][idx + 1] <= j
             ):
                 children.append(make_str())
 
@@ -91,15 +101,17 @@ def make_str():
 
     return make_str()
 
+def parse_token_string(token):
+    return parse_string(token.doc[token.i : token.i + 1])
 
 def get_subconstituents(span):
     constituent_data, position = get_constituent(span)
-    label_vocab = constituent_data.label_vocab
+    label_vocab = constituent_data["label_vocab"]
     doc = span.doc
 
-    while position < len(constituent_data.starts):
-        start = constituent_data.starts[position]
-        end = constituent_data.ends[position]
+    while position < len(constituent_data["starts"]):
+        start = constituent_data["starts"][position]
+        end = constituent_data["ends"][position]
 
         if span.end <= start or span.end < end:
             break
@@ -110,14 +122,14 @@ def get_subconstituents(span):
 
 def get_child_spans(span):
     constituent_data, position = get_constituent(span)
-    label_vocab = constituent_data.label_vocab
+    label_vocab = constituent_data["label_vocab"]
     doc = span.doc
 
     child_start_expected = span.start
     position += 1
-    while position < len(constituent_data.starts):
-        start = constituent_data.starts[position]
-        end = constituent_data.ends[position]
+    while position < len(constituent_data["starts"]):
+        start = constituent_data["starts"][position]
+        end = constituent_data["ends"][position]
 
         if span.end <= start or span.end < end:
             break
@@ -131,14 +143,14 @@ def get_child_spans(span):
 
 def get_parent_span(span):
     constituent_data, position = get_constituent(span)
-    label_vocab = constituent_data.label_vocab
+    label_vocab = constituent_data["label_vocab"]
     doc = span.doc
     sent = span.sent
 
     position -= 1
     while position >= 0:
-        start = constituent_data.starts[position]
-        end = constituent_data.ends[position]
+        start = constituent_data["starts"][position]
+        end = constituent_data["ends"][position]
 
         if start <= span.start and span.end <= end:
             return doc[start:end]
@@ -148,6 +160,8 @@ def get_parent_span(span):
 
     return None
 
+def get_parent_token(token):
+    return get_parent_span(token.doc[token.i : token.i + 1])
 
 def install_spacy_extensions():
     from spacy.tokens import Doc, Span, Token
@@ -162,14 +176,14 @@ def install_spacy_extensions():
     Span.set_extension("children", getter=get_child_spans)
 
     Token.set_extension(
-        "labels", getter=lambda token: get_labels(token.doc[token.i : token.i + 1])
+        "labels", getter=get_token_labels
     )
     Token.set_extension(
         "parse_string",
-        getter=lambda token: parse_string(token.doc[token.i : token.i + 1]),
+        getter=parse_token_string,
     )
     Token.set_extension(
-        "parent", getter=lambda token: get_parent_span(token.doc[token.i : token.i + 1])
+        "parent", getter=get_parent_token
     )
 
 

src/benepar/integrations/spacy_plugin.py

@@ -25,10 +25,9 @@ def finalize(self, doc, label_vocab):
             if self.starts[position] != prev:
                 prev = self.starts[position]
                 loc_to_constituent[self.starts[position]] = position
-
         return ConstituentData(
             self.starts, self.ends, self.labels, loc_to_constituent, label_vocab
-        )
+        ).serialize()
 
 
 class SentenceWrapper(BaseInputExample):

src/benepar/retokenization.py

@@ -88,7 +88,7 @@ def retokenize(
 class Retokenizer:
     def __init__(self, pretrained_model_name_or_path, retain_start_stop=False):
         self.tokenizer = transformers.AutoTokenizer.from_pretrained(
-            pretrained_model_name_or_path, fast=True
+            pretrained_model_name_or_path, fast=False
         )
         if not self.tokenizer.is_fast:
             raise NotImplementedError(