diff --git a/applications/ColossalChat/coati/reasoner/guided_search/mcts.py b/applications/ColossalChat/coati/reasoner/guided_search/mcts.py index 438d0ed5f..100335ba9 100644 --- a/applications/ColossalChat/coati/reasoner/guided_search/mcts.py +++ b/applications/ColossalChat/coati/reasoner/guided_search/mcts.py @@ -120,14 +120,14 @@ class MCTS(BaseModel): self.back_propagation(child) return self.get_best_answer() - + def _iter_nodes(self): to_visit = deque([self.root]) while to_visit: current_node = to_visit.popleft() yield current_node to_visit.extend(current_node.children) - + def get_best_answer(self): best_node = max(self._iter_nodes(), key=lambda node: node.Q, default=self.root) return best_node.answer