This commit is contained in:
Matthew Honnibal 2019-09-18 13:40:03 +02:00
parent 88a23cf49a
commit fa9a283128

View File

@ -41,7 +41,7 @@ class DocPallet(object):
document from the pallet. document from the pallet.
""" """
def __init__(self, attrs=None, store_user_data=False): def __init__(self, attrs=None, store_user_data=False):
"""Create a DocBox object, to hold serialized annotations. """Create a DocPallet object, to hold serialized annotations.
attrs (list): List of attributes to serialize. 'orth' and 'spacy' are attrs (list): List of attributes to serialize. 'orth' and 'spacy' are
always serialized, so they're not required. Defaults to None. always serialized, so they're not required. Defaults to None.
@ -57,7 +57,7 @@ class DocPallet(object):
self.store_user_data = store_user_data self.store_user_data = store_user_data
def add(self, doc): def add(self, doc):
"""Add a doc's annotations to the DocBox for serialization.""" """Add a doc's annotations to the DocPallet for serialization."""
array = doc.to_array(self.attrs) array = doc.to_array(self.attrs)
if len(array.shape) == 1: if len(array.shape) == 1:
array = array.reshape((array.shape[0], 1)) array = array.reshape((array.shape[0], 1))
@ -86,7 +86,7 @@ class DocPallet(object):
yield doc yield doc
def merge(self, other): def merge(self, other):
"""Extend the annotations of this DocBox with the annotations from another.""" """Extend the annotations of this DocPallet with the annotations from another."""
assert self.attrs == other.attrs assert self.attrs == other.attrs
self.tokens.extend(other.tokens) self.tokens.extend(other.tokens)
self.spaces.extend(other.spaces) self.spaces.extend(other.spaces)
@ -95,7 +95,7 @@ class DocPallet(object):
self.user_data.extend(other.user_data) self.user_data.extend(other.user_data)
def to_bytes(self): def to_bytes(self):
"""Serialize the DocBox's annotations into a byte string.""" """Serialize the DocPallet's annotations into a byte string."""
for tokens in self.tokens: for tokens in self.tokens:
assert len(tokens.shape) == 2, tokens.shape assert len(tokens.shape) == 2, tokens.shape
lengths = [len(tokens) for tokens in self.tokens] lengths = [len(tokens) for tokens in self.tokens]
@ -111,7 +111,7 @@ class DocPallet(object):
return gzip.compress(srsly.msgpack_dumps(msg)) return gzip.compress(srsly.msgpack_dumps(msg))
def from_bytes(self, string): def from_bytes(self, string):
"""Deserialize the DocBox's annotations from a byte string.""" """Deserialize the DocPallet's annotations from a byte string."""
msg = srsly.msgpack_loads(gzip.decompress(string)) msg = srsly.msgpack_loads(gzip.decompress(string))
self.attrs = msg["attrs"] self.attrs = msg["attrs"]
self.strings = set(msg["strings"]) self.strings = set(msg["strings"])
@ -134,7 +134,7 @@ def merge_boxes(boxes):
merged = None merged = None
for byte_string in boxes: for byte_string in boxes:
if byte_string is not None: if byte_string is not None:
box = DocBox(store_user_data=True).from_bytes(byte_string) box = DocPallet(store_user_data=True).from_bytes(byte_string)
if merged is None: if merged is None:
merged = box merged = box
else: else:
@ -150,11 +150,11 @@ def pickle_box(box):
def unpickle_box(byte_string): def unpickle_box(byte_string):
return DocBox().from_bytes(byte_string) return DocPallet().from_bytes(byte_string)
copy_reg.pickle(DocBox, pickle_box, unpickle_box) copy_reg.pickle(DocPallet, pickle_box, unpickle_box)
# Compatibility, as we had named it this previously. # Compatibility, as we had named it this previously.
Binder = DocBox Binder = DocPallet
__all__ = ["DocPallet"] __all__ = ["DocPallet"]