From 3d049af56317b98654bbb0440cfc01c8bb18f5c5 Mon Sep 17 00:00:00 2001 From: Matthew Honnibal Date: Sat, 19 Aug 2017 18:42:11 +0200 Subject: [PATCH] Improve vectors to/from disk --- spacy/vectors.pyx | 117 ++++++++++++---------------------------------- 1 file changed, 31 insertions(+), 86 deletions(-) diff --git a/spacy/vectors.pyx b/spacy/vectors.pyx index 1b1e8000a..ac748292c 100644 --- a/spacy/vectors.pyx +++ b/spacy/vectors.pyx @@ -21,6 +21,7 @@ cdef class Vectors: cdef public object data cdef readonly StringStore strings cdef public object key2row + cdef public object keys def __init__(self, strings, data_or_width): self.strings = StringStore() @@ -31,8 +32,11 @@ cdef class Vectors: data = data_or_width self.data = data self.key2row = {} + self.keys = np.ndarray((self.data.shape[0],), dtype='uint64') for i, string in enumerate(strings): - self.key2row[self.strings.add(string)] = i + key = self.strings.add(string) + self.key2row[key] = i + self.keys[i] = key def __reduce__(self): return (Vectors, (self.strings, self.data)) @@ -70,23 +74,30 @@ cdef class Vectors: raise NotImplementedError def to_disk(self, path, **exclude): - def serialize_vectors(p): - write_vectors_to_bin_loc(self.strings, self.key2row, self.data, str(p)) - serializers = OrderedDict(( - ('vec.bin', serialize_vectors), + ('vectors', lambda p: numpy.save(p.open('wb'), self.data)), + ('strings.json', self.strings.to_disk), + ('keys', lambda p: numpy.save(p.open('wb'), self.keys)), )) - return util.to_disk(serializers, exclude) + return util.to_disk(path, serializers, exclude) def from_disk(self, path, **exclude): - def deserialize_vectors(p): - values = load_vectors_from_bin_loc(self.strings, str(p)) - self.key2row, self.data = values + def load_keys(path): + self.keys = numpy.load(path) + for i, key in enumerate(self.keys): + self.keys[i] = key + self.key2row[key] = i + + def load_vectors(path): + self.data = numpy.load(path) serializers = OrderedDict(( - ('vec.bin', deserialize_vectors), + ('keys', load_keys), + ('vectors', load_vectors), + ('strings.json', self.strings.from_disk), )) - return util.from_disk(path, serializers, exclude) + util.from_disk(path, serializers, exclude) + return self def to_bytes(self, **exclude): def serialize_weights(): @@ -94,9 +105,8 @@ cdef class Vectors: return self.data.to_bytes() else: return msgpack.dumps(self.data) - b = msgpack.dumps(self.key2row) serializers = OrderedDict(( - ('key2row', lambda: msgpack.dumps(self.key2row)), + ('keys', lambda: msgpack.dumps(self.keys)), ('strings', lambda: self.strings.to_bytes()), ('vectors', serialize_weights) )) @@ -109,80 +119,15 @@ cdef class Vectors: else: self.data = msgpack.loads(b) + def load_keys(keys): + for i, key in enumerate(keys): + self.keys[i] = key + self.key2row[key] = i + deserializers = OrderedDict(( - ('key2row', lambda b: self.key2row.update(msgpack.loads(b))), + ('keys', lambda b: load_keys(msgpack.loads(b))), ('strings', lambda b: self.strings.from_bytes(b)), ('vectors', deserialize_weights) )) - return util.from_bytes(deserializers, exclude) - - -def write_vectors_to_bin_loc(StringStore strings, dict key2i, - np.ndarray vectors, out_loc): - - cdef int32_t vec_len = vectors.shape[1] - cdef int32_t word_len - cdef bytes word_str - cdef char* chars - cdef uint64_t key - cdef int32_t i - cdef float* vec - - cdef CFile out_file = CFile(out_loc, 'wb') - keys = [(i, key) for (key, i) in key2i.item()] - keys.sort() - for i, key in keys: - vec = vectors.data[i * vec_len] - word_str = strings[key].encode('utf8') - word_len = len(word_str) - - out_file.write_from(&word_len, 1, sizeof(word_len)) - out_file.write_from(&vec_len, 1, sizeof(vec_len)) - - chars = word_str - out_file.write_from(chars, word_len, sizeof(char)) - out_file.write_from(vec, vec_len, sizeof(float)) - out_file.close() - - -def load_vectors_from_bin_loc(StringStore strings, loc): - """ - Load vectors from the location of a binary file. - Arguments: - loc (unicode): The path of the binary file to load from. - Returns: - vec_len (int): The length of the vectors loaded. - """ - cdef CFile file_ = CFile(loc, b'rb') - cdef int32_t word_len - cdef int32_t vec_len = 0 - cdef int32_t prev_vec_len = 0 - cdef float* vec - cdef attr_t string_id - cdef bytes py_word - cdef vector[float*] vectors - cdef int line_num = 0 - cdef Pool mem = Pool() - cdef dict key2i = {} - while True: - try: - file_.read_into(&word_len, sizeof(word_len), 1) - except IOError: - break - file_.read_into(&vec_len, sizeof(vec_len), 1) - if prev_vec_len != 0 and vec_len != prev_vec_len: - raise Exception("Mismatched vector sizes") - if 0 >= vec_len >= MAX_VEC_SIZE: - raise Exception("Mismatched vector sizes") - - chars = file_.alloc_read(mem, word_len, sizeof(char)) - vec = file_.alloc_read(mem, vec_len, sizeof(float)) - - key = strings.add(chars[:word_len]) - key2i[key] = vectors.size() - vectors.push_back(vec) - numpy_vectors = numpy.zeros((vectors.size(), vec_len), dtype='f') - for i in range(vectors.size()): - for j in range(vec_len): - numpy_vectors[i, j] = vectors[i][j] - return key2i, numpy_vectors + util.from_bytes(deserializers, exclude) + return self