diff --git a/data/omegaPRM_v2/omegaprm.py b/data/omegaPRM_v2/omegaprm.py index e1c2d70..263fd0f 100644 --- a/data/omegaPRM_v2/omegaprm.py +++ b/data/omegaPRM_v2/omegaprm.py @@ -6,7 +6,7 @@ from typing import List, Tuple, Dict, Any, Optional import itertools from llm_utils import LLMService - +self.visited_rollouts = [] # Helper function to separate reasoning steps def separate_steps(steps: List[str], mode: str = 'join') -> Any: delimiter = "\n\n" @@ -109,6 +109,7 @@ def __init__(self, solution_prefix: str, parent: Optional['State'] = None): self.R: List[str] = [] # Set of all rollouts from this state self.incorrect_rollouts: List[str] = [] # List of incorrect rollouts self.children: List['State'] = [] # List of child states + self.visited_rollouts = [] def add_rollout(self, rollout: str): self.R.append(rollout) @@ -178,6 +179,8 @@ def add_or_update(self, state: State, rollout: str, priority: float): - rollout (str): The rollout string. - priority (float): The new priority score. """ + if rollout in state.visited_rollouts: + return state_id = id(state) # Unique identifier for the state object rollout_key = (state_id, rollout) @@ -301,7 +304,7 @@ def run(self, question: str, answer: str) -> List: if selected_state is None or selected_rollout is None: # print("No more candidates to explore. Terminating search.\n") break - + selected_state.visited_rollouts.append(selected_rollout) self.expansion_phase_binary_search(selected_state, selected_rollout) # Maintenance Phase