softmax.py 文件源码

def build_huffman_tree(freqs):
    nodes = [HierarchicalSoftmaxNode() for _ in freqs]
    heap = list(zip(freqs, nodes))
    heapq.heapify(heap)

    def pop_initialize(parent, branch):
        freq, node = heapq.heappop(heap)
        node.parent = parent
        node.branch = branch
        return freq

    idx = len(nodes) - 1
    while len(heap) > 1:
        idx += 1
        node = HierarchicalSoftmaxNode()
        nodes.append(node)
        freq = pop_initialize(idx, True) + pop_initialize(idx, False)
        heapq.heappush(heap, (freq, node))
    assert len(heap) == 1
    return nodes