spaCy/spacy/tokens/_serialize.py
Ines Montani f37863093a 💫 Replace ujson, msgpack and dill/pickle/cloudpickle with srsly (#3003)
Remove hacks and wrappers, keep code in sync across our libraries and move spaCy a few steps closer to only depending on packages with binary wheels 🎉

See here: https://github.com/explosion/srsly

    Serialization is hard, especially across Python versions and multiple platforms. After dealing with many subtle bugs over the years (encodings, locales, large files) our libraries like spaCy and Prodigy have steadily grown a number of utility functions to wrap the multiple serialization formats we need to support (especially json, msgpack and pickle). These wrapping functions ended up duplicated across our codebases, so we wanted to put them in one place.

    At the same time, we noticed that having a lot of small dependencies was making maintainence harder, and making installation slower. To solve this, we've made srsly standalone, by including the component packages directly within it. This way we can provide all the serialization utilities we need in a single binary wheel.

    srsly currently includes forks of the following packages:

        ujson
        msgpack
        msgpack-numpy
        cloudpickle



* WIP: replace json/ujson with srsly

* Replace ujson in examples

Use regular json instead of srsly to make code easier to read and follow

* Update requirements

* Fix imports

* Fix typos

* Replace msgpack with srsly

* Fix warning
2018-12-03 01:28:22 +01:00

118 lines
4.1 KiB
Python

from __future__ import unicode_literals
import numpy
import gzip
import srsly
from thinc.neural.ops import NumpyOps
from ..compat import copy_reg
from ..tokens import Doc
from ..attrs import SPACY, ORTH
class Binder(object):
"""Serialize analyses from a collection of doc objects."""
def __init__(self, attrs=None):
"""Create a Binder object, to hold serialized annotations.
attrs (list):
List of attributes to serialize. 'orth' and 'spacy' are always
serialized, so they're not required. Defaults to None.
"""
attrs = attrs or []
self.attrs = list(attrs)
# Ensure ORTH is always attrs[0]
if ORTH in self.attrs:
self.attrs.pop(ORTH)
if SPACY in self.attrs:
self.attrs.pop(SPACY)
self.attrs.insert(0, ORTH)
self.tokens = []
self.spaces = []
self.strings = set()
def add(self, doc):
"""Add a doc's annotations to the binder for serialization."""
array = doc.to_array(self.attrs)
if len(array.shape) == 1:
array = array.reshape((array.shape[0], 1))
self.tokens.append(array)
spaces = doc.to_array(SPACY)
assert array.shape[0] == spaces.shape[0]
spaces = spaces.reshape((spaces.shape[0], 1))
self.spaces.append(numpy.asarray(spaces, dtype=bool))
self.strings.update(w.text for w in doc)
def get_docs(self, vocab):
"""Recover Doc objects from the annotations, using the given vocab."""
for string in self.strings:
vocab[string]
orth_col = self.attrs.index(ORTH)
for tokens, spaces in zip(self.tokens, self.spaces):
words = [vocab.strings[orth] for orth in tokens[:, orth_col]]
doc = Doc(vocab, words=words, spaces=spaces)
doc = doc.from_array(self.attrs, tokens)
yield doc
def merge(self, other):
"""Extend the annotations of this binder with the annotations from another."""
assert self.attrs == other.attrs
self.tokens.extend(other.tokens)
self.spaces.extend(other.spaces)
self.strings.update(other.strings)
def to_bytes(self):
"""Serialize the binder's annotations into a byte string."""
for tokens in self.tokens:
assert len(tokens.shape) == 2, tokens.shape
lengths = [len(tokens) for tokens in self.tokens]
msg = {
"attrs": self.attrs,
"tokens": numpy.vstack(self.tokens).tobytes("C"),
"spaces": numpy.vstack(self.spaces).tobytes("C"),
"lengths": numpy.asarray(lengths, dtype="int32").tobytes("C"),
"strings": list(self.strings),
}
return gzip.compress(srsly.msgpack_dumps(msg))
def from_bytes(self, string):
"""Deserialize the binder's annotations from a byte string."""
msg = srsly.msgpack_loads(gzip.decompress(string))
self.attrs = msg["attrs"]
self.strings = set(msg["strings"])
lengths = numpy.fromstring(msg["lengths"], dtype="int32")
flat_spaces = numpy.fromstring(msg["spaces"], dtype=bool)
flat_tokens = numpy.fromstring(msg["tokens"], dtype="uint64")
shape = (flat_tokens.size // len(self.attrs), len(self.attrs))
flat_tokens = flat_tokens.reshape(shape)
flat_spaces = flat_spaces.reshape((flat_spaces.size, 1))
self.tokens = NumpyOps().unflatten(flat_tokens, lengths)
self.spaces = NumpyOps().unflatten(flat_spaces, lengths)
for tokens in self.tokens:
assert len(tokens.shape) == 2, tokens.shape
return self
def merge_bytes(binder_strings):
"""Concatenate multiple serialized binders into one byte string."""
output = None
for byte_string in binder_strings:
binder = Binder().from_bytes(byte_string)
if output is None:
output = binder
else:
output.merge(binder)
return output.to_bytes()
def pickle_binder(binder):
return (unpickle_binder, (binder.to_bytes(),))
def unpickle_binder(byte_string):
return Binder().from_bytes(byte_string)
copy_reg.pickle(Binder, pickle_binder, unpickle_binder)