This function is used to extract solution from a built tree. It is mainly used for MCTS, but also works for saved tree from step_beam.
(
full_tree_dict: Dict[str, Any],
prune: bool = False,
b1: int = 1,
b2: int = 5,
strategy: str = "q_value",
c_puct: float = 1.25,
)
| 128 | raise NotImplementedError(f"strategy {strategy} not implemented") |
| 129 | |
| 130 | def get_solution( |
| 131 | full_tree_dict: Dict[str, Any], |
| 132 | prune: bool = False, |
| 133 | b1: int = 1, |
| 134 | b2: int = 5, |
| 135 | strategy: str = "q_value", |
| 136 | c_puct: float = 1.25, |
| 137 | ) -> Optional[Dict[str, Any]]: |
| 138 | """ |
| 139 | This function is used to extract solution from a built tree. |
| 140 | It is mainly used for MCTS, but also works for saved tree from step_beam. |
| 141 | """ |
| 142 | question = full_tree_dict["question"] |
| 143 | ground_truth = full_tree_dict.get("answer", None) |
| 144 | tree_dict = full_tree_dict["react"] |
| 145 | |
| 146 | # rebuild tree |
| 147 | root, tree_depth = rebuild_tree(tree_dict, max_num_children=b1*b2, c_puct=c_puct) |
| 148 | |
| 149 | # pruning tree |
| 150 | if prune: |
| 151 | prune_node(root) |
| 152 | if root.prune: |
| 153 | # no valid leaf node for the entire tree |
| 154 | return None |
| 155 | |
| 156 | # search in tree |
| 157 | final_answer_nodes = [] |
| 158 | current_top_num = b1 |
| 159 | current_nodes = [root] |
| 160 | |
| 161 | for _ in range(tree_depth): |
| 162 | candidate_nodes = select_non_prune(current_nodes) |
| 163 | candidate_nodes = sort_by_strategy(candidate_nodes, strategy) |
| 164 | current_nodes = candidate_nodes[:current_top_num] |
| 165 | |
| 166 | for current_node in current_nodes[:]: |
| 167 | if is_valid_final_answer_node(current_node): |
| 168 | final_answer_nodes.append(current_node) |
| 169 | current_nodes.remove(current_node) |
| 170 | current_top_num -= 1 |
| 171 | elif not current_node.children: |
| 172 | current_nodes.remove(current_node) |
| 173 | current_top_num -= 1 |
| 174 | |
| 175 | if not final_answer_nodes: |
| 176 | return None |
| 177 | |
| 178 | final_answer_nodes = sort_by_strategy(final_answer_nodes, strategy) |
| 179 | top_final_answer_node = final_answer_nodes[0] |
| 180 | |
| 181 | # for node in final_answer_nodes: |
| 182 | # print(node.tag) |
| 183 | |
| 184 | return { |
| 185 | "question": question, |
| 186 | "ground_truth": ground_truth, |
| 187 | "final_answer": top_final_answer_node.final_answer, |
no test coverage detected