* Work on huffman coding.

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 | ((<uint32_t>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 = - <int64_t>id_
+    leaf_id = - <int64_t>(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=- <int64_t>id1, right=- <int64_t>id2, prob=prob))
+cdef int append_two(vector[Node]& nodes, int id1, int id2, float prob) except -1:
+    cdef Node node
+    node.left = -<int64_t>(id1 + 1)
+    node.right = -<int64_t>(id2 + 1)
+    node.prob = prob
+    nodes.push_back(node)