diff --git a/spacy/serialize.pyx b/spacy/serialize.pyx index bbfce1334..8efa049da 100644 --- a/spacy/serialize.pyx +++ b/spacy/serialize.pyx @@ -1,52 +1,122 @@ -cdef class HuffmanTable: - def __init__(self, Vocab vocab): - cdef int i = vocab.lexemes.size()-1 - j = 0 - cdef vector[Node] nodes - while i >= 0 and j < size: - if i >= 1 and lexemes[i-1].prob < nodes[j].prob: - append_two(&nodes, j, lexemes[i].id, lexemes[i+1].id, - lexemes[i].prob + lexemes[i+1].prob) - i -= 2 - elif (j + 1) < size and nodes[j+1].prob < lexemes[i].prob: - append_zero(&nodes, j) - j += 2 - else: - append_one(&nodes, j, lexemes[i].id, lexemes[i].prob) - i -= 1 - j += 1 - assign_codes(&nodes, j, 0, 0) - i = 0 - for i in range(nodes.length): - node = nodes[i] - if node.left < 0: - vocab.codes[- node.left] = node.code - if nodes[i].right < 0: - vocab.codes[- node.right] = set_another_bit(node.code) +from libcpp.vector cimport vector +from libc.stdint cimport uint32_t +from libc.stdint cimport int64_t +from libc.stdint cimport uint64_t + +import numpy -cdef int assign_codes(vector[Node]* nodes, int i, uint32_t code, uint32_t j) except -1: - nodes[i].code = code +cdef struct Node: + float prob + int left + int right + + +cdef struct BitArray: + uint64_t data + + +cdef BitArray set_bit(BitArray barray, unsigned char id_) nogil: + cdef uint64_t one = 1 + barray.data |= one << id_ + return barray + + +cdef BitArray clear_bit(BitArray barray, unsigned char id_) nogil: + cdef uint64_t one = 1 + barray.data ^= one << id_ + return barray + + +cpdef uint64_t[:] huffman_encode(float[:] probs): + assert len(probs) >= 3 + + output = numpy.zeros(shape=(len(probs),), dtype=numpy.uint64) + + cdef int size = len(probs) + cdef vector[Node] nodes + cdef int i = size - 1 + cdef int j = 0 + + append_two(nodes, i, i-1, probs[i] + probs[i-1]) + i -= 2 + append_one(nodes, 0, i, probs[i]) + j += 1 + i -= 1 + while i >= 0 or j < len(nodes): + if i < 0: + append_zero(nodes, j) + j += 2 + elif j >= len(nodes): + append_two(nodes, i, i-1, probs[i]+probs[i-1]) + elif i >= 1 and (j == len(nodes) or probs[i-1] < nodes[j].prob): + append_two(nodes, i, i-1, probs[i] + probs[i-1]) + i -= 2 + elif (j+1) < len(nodes) and nodes[j+1].prob < probs[i]: + append_zero(nodes, j) + j += 2 + else: + append_one(nodes, j, i, probs[i]) + i -= 1 + j += 1 + cdef vector[BitArray] codes + codes.resize(len(probs)) + for i in range(len(probs)): + codes[i].data = 0 + assign_codes(nodes, codes, len(nodes) - 2, BitArray(data=0), 0) + output = numpy.zeros(shape=(len(codes),), dtype=numpy.uint64) + for i in range(len(codes)): + output[i] = codes[i].data + return output + + +cdef int assign_codes(vector[Node]& nodes, vector[BitArray]& codes, int i, + BitArray code, int bit) except -1: + cdef BitArray left_code = clear_bit(code, bit) if nodes[i].left >= 0: - assign_codes(nodes, nodes[i].left, code, j+1) + if nodes[i].left != i: + assign_codes(nodes, codes, nodes[i].left, left_code, bit+1) + else: + id_ = -(nodes[i].left + 1) + codes[id_] = left_code + cdef BitArray right_code = set_bit(code, bit) if nodes[i].right >= 0: - assign_codes(nodes, nodes[i].right, code | ((1) << j), j+1) + if nodes[i].right != i: + assign_codes(nodes, codes, nodes[i].right, right_code, bit+1) + else: + id_ = -(nodes[i].right + 1) + codes[id_] = right_code -cdef int append_zero(vector[Node]* nodes, int j) except -1: - nodes.push_back(Node(left=j, right=j+1, prob=nodes[j].prob+nodes[j+1].prob)) +cdef int append_zero(vector[Node]& nodes, int j) nogil: + cdef Node node + node.left = j + node.right = j+1 + node.prob = nodes[j].prob + nodes[j+1].prob + nodes.push_back(node) -cdef int append_one(vector[Node]* nodes, int j, attr_t id_, weight_t prob) except -1: +cdef int append_one(vector[Node]& nodes, int j, int id_, float prob) except -1: + cdef Node node # Encode leaves as negative integers, where the integer is the index of the # word in the vocabulary. - leaf_id = - id_ + leaf_id = - (id_ + 1) new_prob = prob + nodes[j].prob if prob < nodes[j].prob: - nodes.push_back(Node(left=leaf_id, right=j, prob=new_prob)) + node.left = leaf_id + node.right = j + node.prob = new_prob + nodes.push_back(node) else: - nodes.push_back(Node(left=j, right=leaf_id, prob=new_prob)) + node.left = j + node.right = leaf_id + node.prob = new_prob + nodes.push_back(node) -cdef int append_two(vector[Node]* nodes, attr_t id1, attr_t id2, weight_t prob) except -1: - nodes.push_back(Node(left=- id1, right=- id2, prob=prob)) +cdef int append_two(vector[Node]& nodes, int id1, int id2, float prob) except -1: + cdef Node node + node.left = -(id1 + 1) + node.right = -(id2 + 1) + node.prob = prob + nodes.push_back(node)