| 62 | } |
| 63 | |
| 64 | DDTree build_ddtree(const float * top_log_probs, |
| 65 | const int32_t * top_token_ids, |
| 66 | int L, int K, int budget, |
| 67 | bool chain_seed) { |
| 68 | DDTree tree; |
| 69 | if (budget <= 0 || L <= 0) { |
| 70 | tree.parents.push_back(-1); |
| 71 | tree.child_maps.emplace_back(); |
| 72 | tree.visibility.assign(1, 1); |
| 73 | return tree; |
| 74 | } |
| 75 | |
| 76 | struct HeapEntry { |
| 77 | float neg_logw; |
| 78 | std::vector<int> ranks; |
| 79 | int parent_index; |
| 80 | int depth; |
| 81 | int rank; |
| 82 | float logw; |
| 83 | }; |
| 84 | struct HeapCmp { |
| 85 | bool operator()(const HeapEntry & a, const HeapEntry & b) const { |
| 86 | return a.neg_logw > b.neg_logw; |
| 87 | } |
| 88 | }; |
| 89 | std::priority_queue<HeapEntry, std::vector<HeapEntry>, HeapCmp> heap; |
| 90 | |
| 91 | tree.token_ids.reserve(budget); |
| 92 | tree.depths.reserve(budget); |
| 93 | tree.parents.reserve(budget + 1); |
| 94 | tree.parents.push_back(-1); |
| 95 | tree.child_maps.emplace_back(); |
| 96 | |
| 97 | if (chain_seed) { |
| 98 | const int chain_depth = std::min(L, budget); |
| 99 | float cum_logw = 0.0f; |
| 100 | int prev_idx = 0; |
| 101 | for (int d = 1; d <= chain_depth; d++) { |
| 102 | const int32_t tok_id = top_token_ids[(size_t)(d - 1) * K + 0]; |
| 103 | cum_logw += top_log_probs[(size_t)(d - 1) * K + 0]; |
| 104 | |
| 105 | const int cur_idx = tree.n_nodes + 1; |
| 106 | tree.token_ids.push_back(tok_id); |
| 107 | tree.depths.push_back(d); |
| 108 | tree.parents.push_back(prev_idx); |
| 109 | tree.child_maps.emplace_back(); |
| 110 | tree.child_maps[prev_idx][tok_id] = cur_idx; |
| 111 | tree.n_nodes++; |
| 112 | |
| 113 | if (K > 1) { |
| 114 | const float sibling_logw = cum_logw |
| 115 | - top_log_probs[(size_t)(d - 1) * K + 0] |
| 116 | + top_log_probs[(size_t)(d - 1) * K + 1]; |
| 117 | heap.push({ |
| 118 | -sibling_logw, |
| 119 | {1}, |
| 120 | prev_idx, |
| 121 | d, |