github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge branch 'feat/prm' of github.com:TongLi3701/ColossalAI into feat/prm

pull/6119/head
Tong Li 2024-11-08 12:43:07 +00:00
parent ed817a29f9 308960534f
commit dcb509c8e3
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
applications/ColossalChat/coati/reasoner/guided_search/mcts.py
View File

@ -119,14 +119,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