omegaprm.py

import heapq
import math
import random
import re
import json
from typing import List, Tuple, Dict, Any, Optional
import itertools
from transformers import AutoTokenizer
import asyncio  # New import added for async handling
from openai import AsyncOpenAI   # Using AsyncOpenAI as client
import numpy as np


# Helper function to separate reasoning steps
def separate_steps(steps: List[str], mode: str = 'join') -> Any:
    delimiter = "\n\n"
    if mode == 'join':
        if not isinstance(steps, list):
            raise TypeError("For 'join' mode, 'steps' must be a list of strings.")
        return delimiter.join(steps)
    elif mode == 'split':
        if not isinstance(steps, str):
            raise TypeError("For 'split' mode, 'steps' must be a string.")
        return steps.split(delimiter)
    else:
        raise ValueError("Mode should be either 'join' or 'split'.")
    

# def judge_ans(
#         problem_str: str,
#         extracted_groundtruth: str,
#         output_list: List[str],
#         v_list: List[float],
#         aggration_mode: str,
#         extract_answer_fn,
#         judge_correct_fn,
#         normalize=False,
#     ):
#         ans_list = [extract_answer_fn(txt) for txt in output_list]
#         valid_ans_list, valid_v_list = [], []
#         for i, ans in enumerate(ans_list):
#             if ans != INVALID_ANS:
#                 valid_ans_list.append(ans)
#                 valid_v_list.append(v_list[i])
#         if len(valid_ans_list) == 0:
#             return 0

#         if "orm" in aggration_mode and normalize:
#             # score_normalization: this is only necessary for [-1, 1] values
#             valid_v_list = np.array(valid_v_list)
#             valid_v_list -= valid_v_list.min()
#             valid_v_list /= valid_v_list.max() + 1e-3
#             valid_v_list = valid_v_list.tolist()
#         aggregated_ans = AGG_FN_MAP[aggration_mode](valid_ans_list, valid_v_list)

#         return (
#             1 if judge_correct_fn(problem_str, extracted_groundtruth, aggregated_ans) else 0
#         )





# Helper function to check correctness of a generated response
def check_correctness(generated_response: str, expected_answer: str) -> bool:
    # sentences = re.split(
    #     r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s', generated_response.strip()
    # )
    # last_sentence = sentences[-1] if sentences else ''
    # return expected_answer.strip() in last_sentence.strip()
    extract_answer_fn = MATH.extract_answer
    judge_correct_fn = MATH.judge_correct
    answer = extract_answer_fn(generated_response)
    return (
            1 if judge_correct_fn("", expected_answer, answer) else 0
        )


class ProcessRewardModel:
    """
    ProcessRewardModel encapsulates the reward inference process.
    
    It utilizes a chat-based reward model (e.g., 'Llama3.1-8B-PRM-Mistral-Data')
    to evaluate a sequence of reasoning steps. It iteratively sends messages to the model,
    checking that each step receives a positive judgement (i.e., its completion starts with '+').
    """

    def __init__(self, client, model="Llama3.1-8B-PRM-Mistral-Data", temperature=0.0, max_tokens=1):
        """
        Initialize the ProcessRewardModel.

        Parameters:
            client: The chat client instance that provides a `chat.completions.create` method.
            model (str): The model name to be used for generating reward completions.
            temperature (float): Sampling temperature. Default is 0.0 for deterministic outcomes.
            max_tokens (int): Maximum tokens to generate in the reward inference. Default is 1.
        """
        self.client = client
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens

    def evaluate(self, problem: str, steps: list, output_type: str = 'bool') -> bool:
        """
        Synchronously evaluate the process reward using asynchronous API calls.
        This method wraps the asynchronous _async_evaluate call.

        Parameters:
            problem (str): The problem or question statement.
            steps (List[str]): A list of reasoning steps.

        Returns:
            bool: True if all steps are positively judged, False otherwise.
        """
        return asyncio.run(self._async_evaluate(problem, steps, output_type))

    async def _async_evaluate(self, problem: str, steps: list, output_type: str = 'bool') -> bool:
        messages = []

        # Merge every 5 steps into 1 step to reduce evaluation time
        if 'deepseek' in self.model.lower():
            merged_steps = []
            current_merge = []
            for step in steps:
                current_merge.append(step)
                if len(current_merge) == 6:
                    merged_steps.append("\n\n".join(current_merge))
                    current_merge = []
            if current_merge:  # Add any remaining steps
                merged_steps.append("\n\n".join(current_merge))

        steps = merged_steps
        for sdx, step in enumerate(steps):
            if sdx == 0:
                messages.append({
                    'role': 'user',
                    'content': f"{problem}\n\n{step}"
                })
            else:
                messages.append({
                    'role': 'user',
                    'content': step
                })

            completion = await self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                n=1,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
            )

            response = completion.choices[0].message.content.strip().lower()
            if not response.startswith('+'):
                if output_type == 'bool':
                    return False
                else:
                    return [1.0 if i < sdx else 0.0 for i in range(len(steps))]
            messages.append({'role': 'assistant', 'content': '+'})
        
        if output_type == 'bool':
            return True
        else:
            return [1.0] * len(steps)
        
    async def _async_evaluate_system2(self, problem: str, steps: list, output_type: str = 'bool') -> bool:
        messages = []

        # Merge every 5 steps into 1 step to reduce evaluation time

        if 'deepseek' in self.model.lower():
            merged_steps = []
            current_merge = []
            for step in steps:
                current_merge.append(step)
                if len(current_merge) == 6:
                    merged_steps.append("\n\n".join(current_merge))
                    current_merge = []
            if current_merge:  # Add any remaining steps
                merged_steps.append("\n\n".join(current_merge))

        steps = merged_steps
        for sdx, step in enumerate(steps):
            
            if sdx == 0:
                messages.append({
                    'role': 'user', 
                    'content': f"Problem: {problem}\n\nStep: {step}\n\nIs this step correct? You must answer with '+' for correct or '-' for incorrect in the end of your response."
                })
            else:
                messages.append({
                    'role': 'user', 
                    'content': f"Step: {step}\n\nIs this step correct? You must answer with '+' for correct or -' for incorrect in the end of your response."
                })
            
            completion = await self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                n=1,
                temperature=self.temperature,
                max_tokens=8192,
            )
            response = completion.choices[0].message.content

            # print("DyVer Verification:", response)
            
            # New negative checking logic
            content = response.strip().lower()
            last_words = ' '.join(content.split()[-3:])  # Last 3 words
            
            judgment = any(
                '+' in part and '-' not in part
                for part in (
                    content[-5:], 
                    last_words,
                )
            )
            
            if not judgment:
                return [1.0 if i < sdx else 0.0 for i in range(len(steps))]
            messages.append({'role': 'assistant', 'content': '<think>\n\n</think> +'})
        return [1.0] * len(steps)



class LanguageModel:
    def __init__(self, client, model_name="/root/.cache/modelscope/hub/Qwen/Qwen2___5-Math-7B-Instruct",
                 max_new_tokens=512, temperature=0.7, top_p=0.9):
        """
         Initialize the LanguageModel for async OpenAI calls.
         Removed the LLMService dependency and using async calls via openai.
         
         Parameters:
         - client: An instance of AsyncOpenAI passed externally.
         - model_name (str): API model name to use.
         - max_new_tokens (int): Maximum tokens for generation.
         - temperature (float): Sampling temperature.
         - top_p (float): Nucleus sampling probability.
         """
        self.model_name = model_name
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.default_prompt = (
            "Please complete the answer for the question based on the given steps without generating existing steps again, "
            "and separate your following steps using \n\n.\n\n"
        )
        # Retain tokenizer for chat template operations elsewhere.
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(f"/root/{model_name}")
        except Exception as e:
            print(f"Error loading tokenizer: {e}")
            self.tokenizer = None
        # Use the external AsyncOpenAI client.
        self.async_client = client

    async def generate_rollout(self, state_prefix: str, num_copies: int) -> List[str]:
        """
        Asynchronously generate responses using OpenAI's ChatCompletion API.
        
        Parameters:
        - state_prefix (str): The current solution prefix.
        - num_copies (int): The number of response copies to generate.
        
        Returns:
        - List[str]: A list of generated responses.
        """
        response = await self.async_client.completions.create(
            model=self.model_name,
            prompt=state_prefix,
            max_tokens=self.max_new_tokens,
            temperature=self.temperature,
            top_p=self.top_p,
            n=num_copies,
        )
        return [choice.text for choice in response.choices]

    def update_prompt(self, new_prompt: str):
        """
        Update the default prompt if necessary.
        
        Parameters:
        - new_prompt (str): The new prompt template.
        """
        self.default_prompt = new_prompt

    def evaluate_correctness(self, response: str, expected_answer: str) -> bool:
        """
        Check if the generated solution matches the expected answer.
        
        Parameters:
        - response (str): The complete generated response.
        - expected_answer (str): The expected answer to compare with.
        
        Returns:
        - bool: True if the expected answer is in the final part of the solution.
        """
        return check_correctness(response, expected_answer)


# Define the State class
class State:
    def __init__(self, solution_prefix: str, parent: Optional['State'] = None):
        self.solution_prefix = solution_prefix  # Solution prefix as a single string
        self.parent = parent  # Reference to the parent state
        self.N = 0  # Visit count (number of times selected)
        self.total_rollouts = 0  # Total number of rollouts generated from this state
        self.correct_rollouts = 0  # Number of correct rollouts
        self.MC: Optional[float] = None  # Monte Carlo estimation (c/k)
        self.Q: Dict[str, float] = {}  # Q(s, r): estimated value for each rollout
        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

    def add_rollout(self, rollout: str):
        self.R.append(rollout)

    def add_incorrect_rollout(self, rollout: str):
        if rollout not in self.incorrect_rollouts:
            self.incorrect_rollouts.append(rollout)

    def get_full_solution(self) -> str:
        # Return the complete solution from the root to this state
        if self.parent:
            return self.parent.get_full_solution() + '\n\n' + self.solution_prefix
        else:
            return self.solution_prefix

    def get_new_text(self) -> str:
        """
        Return the new text added at this node compared to the parent.
        """
        if self.parent:
            parent_text = self.parent.solution_prefix
            new_text = self.solution_prefix[len(parent_text):].strip()
            return new_text
        else:
            # Root node (the question)
            return self.solution_prefix.strip()

    def get_text_with_labels(self) -> Dict[str, Any]:
        """
        Return a nested dictionary where each node contains:
        - 'text': The new text at this node.
        - 'mc_value': The MC value at this node.
        - 'children': A list of child nodes with the same structure.
        """
        data = {
            'text': self.get_new_text(),
            'mc_value': self.MC,
            'children': [child.get_text_with_labels() for child in self.children]
        }
        return data


# Define the Search Tree class
class SearchTree:
    def __init__(self):
        self.root: Optional[State] = None
        self.nodes: List[State] = []  # List of all states

    def add_state(self, state: State):
        self.nodes.append(state)

# Define the Candidate Pool as a priority queue with update capability
class CandidatePool:
    def __init__(self):
        self.heap: List[Tuple[float, int]] = []  # Heap of (-priority, unique_id)
        self.entry_finder: Dict[int, Tuple[float, int]] = {}  # Maps unique_id to (-priority, unique_id)
        self.counter = itertools.count()  # Unique sequence count
        self.id_to_rollout: Dict[int, Tuple[State, str]] = {}  # Maps unique_id to (state, rollout)
        self.latest_id_per_rollout: Dict[Tuple[int, str], int] = {}  # Maps (state_id, rollout) to unique_id

    def add_or_update(self, state: State, rollout: str, priority: float):
        """
        Add a new rollout or update the priority of an existing rollout.

        Parameters:
        - state (State): The state associated with the rollout.
        - rollout (str): The rollout string.
        - priority (float): The new priority score.
        """
        state_id = id(state)  # Unique identifier for the state object
        rollout_key = (state_id, rollout)

        # Check if the rollout already exists in the pool
        if rollout_key in self.latest_id_per_rollout:
            # Previous unique_id exists; it is now outdated
            old_unique_id = self.latest_id_per_rollout[rollout_key]
            # Mark the old entry as invalid by removing it from entry_finder
            if old_unique_id in self.entry_finder:
                del self.entry_finder[old_unique_id]
                del self.id_to_rollout[old_unique_id]

        # Assign a new unique_id for the updated rollout
        unique_id = next(self.counter)
        self.latest_id_per_rollout[rollout_key] = unique_id

        # Add the new entry to the heap and mappings
        heapq.heappush(self.heap, (-priority, unique_id))  # Max-heap using negative priority
        self.entry_finder[unique_id] = (-priority, unique_id)
        self.id_to_rollout[unique_id] = (state, rollout)

    def pop(self) -> Tuple[Optional[State], Optional[str]]:
        """
        Pop the rollout with the highest priority.

        Returns:
        - Tuple[Optional[State], Optional[str]]: The state and rollout string, or (None, None) if empty.
        """
        while self.heap:
            neg_priority, unique_id = heapq.heappop(self.heap)
            # Check if this unique_id is still valid
            if unique_id in self.entry_finder:
                # Valid entry
                state, rollout = self.id_to_rollout.pop(unique_id)
                del self.entry_finder[unique_id]
                # Remove from latest_id_per_rollout
                state_id = id(state)
                rollout_key = (state_id, rollout)
                if self.latest_id_per_rollout.get(rollout_key) == unique_id:
                    del self.latest_id_per_rollout[rollout_key]
                return state, rollout
            # Else, outdated entry; skip
        return None, None

    def is_empty(self) -> bool:
        return not self.entry_finder

# Define the OmegaPRM algorithm
class OmegaPRM:
    def __init__(self, LM: LanguageModel, reward_model, c_puct: float, alpha: float, beta: float, L: int, k: int, N: int,
                 rollout_budget: int, save_data_tree: bool):
        """
        Initialize the OmegaPRM algorithm.

        Parameters:
            LM (LanguageModel): The language model instance.
            reward_model: An instance of ProcessRewardModel to evaluate solution correctness.
            c_puct (float): Exploration constant.
            alpha (float): Weight for MC(s).
            beta (float): Length penalty.
            L (int): Maximum solution length.
            k (int): Number of rollouts for Monte Carlo estimation.
            N (int): Maximum search count.
            rollout_budget (int): Total rollout budget.
            save_data_tree (bool): Whether to save and return the data tree.
        """
        self.LM = LM
        self.reward_model = reward_model
        self.expected_answer = None
        self.c_puct = c_puct
        self.alpha = alpha
        self.beta = beta
        self.L = L
        self.k = k
        self.N = N
        self.rollout_budget = rollout_budget
        self.save_data_tree = save_data_tree

        self.T = SearchTree()
        self.C = CandidatePool()

        self.n = 0
        self.total_rollouts = 0

    def reset(self):
        """Reset internal state variables to prepare for a fresh run."""
        self.expected_answer = None
        self.T = SearchTree()  # Reset search tree
        self.C = CandidatePool()  # Reset candidate pool
        self.n = 0
        self.total_rollouts = 0
        self.collected_data = []  # Clear collected data

    async def monte_carlo_estimation(self, state: State):
        """
        Perform Monte Carlo estimation for state by generating k rollouts
        and computing MC(s) = c / k, where c is the number of correct rollouts.
        """
        c = 0  # Correct rollouts count
        incorrect_rollouts = []
        correct_rollouts = []
        batct_rollouts = await self.LM.generate_rollout(state.solution_prefix, self.k)

        # Increment visit count of selected state
        state.N += 1

        for i, rollout in enumerate(batct_rollouts):
            # Increment number of total rollouts
            self.total_rollouts += 1

            # Generate rollout r_i
            state.add_rollout(rollout)

            # Evaluate correctness of final answer in rollout using the reward model.
            full_solution = (state.solution_prefix + '\n\n' + rollout).strip() if state.solution_prefix else rollout
            steps = separate_steps(full_solution, mode='split')
            # If all steps receive a positive judgment, evaluate returns -1.
            is_correct = await self.reward_model._async_evaluate(self.problem, steps)

            if is_correct:
                c += 1
                correct_rollouts.append(rollout)
            else:
                incorrect_rollouts.append(rollout)
                state.add_incorrect_rollout(rollout)  # Track incorrect rollouts

        # Update total rollouts and correct rollouts
        state.total_rollouts += self.k
        state.correct_rollouts += c
        state.MC = state.correct_rollouts / state.total_rollouts if state.total_rollouts > 0 else 0

        if state.MC == 1.0:
            # Add all correct rollouts to the tree as new states
            for rollout in correct_rollouts:
                self.add_correct_rollout_to_tree(state, rollout)
        elif state.MC == 0.0:
            # State is incorrect; no further action
            for rollout in incorrect_rollouts:
                self.add_incorrect_rollout_to_tree(state, rollout)
            return
        else:
            # 0 < MC(s) < 1.0
            # Add correct rollouts to the tree
            for rollout in correct_rollouts:
                self.add_correct_rollout_to_tree(state, rollout)
            # Add incorrect rollouts to candidate pool with updated priorities
            for rollout in incorrect_rollouts:
                priority = self.compute_selection_score(state, rollout)
                self.C.add_or_update(state, rollout, priority)

    async def run(self, question: str, answer: str) -> List:
        """
        Execute the OmegaPRM algorithm.

        Parameters:
        - question (str): The question to generate solutions for.

        Returns:
        - Collected data: List of dictionaries.
        """
        self.reset()
        self.problem = question  # Store the original question for reward evaluation

        print(f"Running OmegaPRM for question: '{question}'\n")
        # Initialization
        if self.LM.tokenizer is not None:
            question_tamplated = self.LM.tokenizer.apply_chat_template(
                [{"role": "user", "content": question}],
                tokenize=False,
                add_special_tokens=False,
                add_generation_prompt=True
            )
        else:
            question_tamplated = question
        initial_state = State(solution_prefix=question_tamplated, parent=None)
        self.expected_answer = answer
        self.T.root = initial_state
        self.T.add_state(initial_state)
        self.n = 0

        # Monte Carlo Estimation for initial_state
        await self.monte_carlo_estimation(initial_state)

        # Main loop
        while self.n < self.N and self.total_rollouts < self.rollout_budget and not self.C.is_empty():
            # Selection Phase
            selected_state, selected_rollout = self.selection_phase()
            if selected_state is None or selected_rollout is None:
                break

            await self.expansion_phase_binary_search(selected_state, selected_rollout)

            # Maintenance Phase
            self.maintenance_phase(selected_state)

            # Increment search count
            self.n += 1

        if self.save_data_tree:
            data = self.collect_tree_structure()
        else:
            data = self.collect_solution_prefixes()
        return data

    def compute_Q(self, state: State, rollout: str) -> float:
        """
        Compute Q(s, r) = alpha^{1 - MC(s)} * beta^{len(r)/L}, where len(r) is based on word count.
        """
        # Count words in the rollout
        word_count = len(rollout.split())
        length_penalty = word_count / self.L
        Q_value = (self.alpha ** (1 - state.MC)) * (self.beta ** length_penalty)
        return Q_value

    def compute_U(self, state: State) -> float:
        """
        Compute U(s) = c_puct * sqrt(sum_{s'} N(s')) / (1 + N(s))
        """
        N_total = sum(s.N for s in self.T.nodes)
        if N_total == 0:
            N_total = 1  # Prevent division by zero
        U_s = self.c_puct * (math.sqrt(N_total)) / (1 + state.N)
        return U_s

    def compute_selection_score(self, state: State, rollout: str) -> float:
        """
        Compute selection score: Score(s, r) = Q(s, r) + U(s)
        """
        Q_s_r = self.compute_Q(state, rollout)
        U_s = self.compute_U(state)
        score = Q_s_r + U_s
        return score

    def selection_phase(self) -> Tuple[Optional[State], Optional[str]]:
        """
        Select (state, rollout) with the highest score from candidate pool C.
        """
        selected_state, selected_rollout = self.C.pop()
        return selected_state, selected_rollout

    def add_correct_rollout_to_tree(self, parent_state: State, rollout: str):
        """
        Add the correct rollout to the tree as a child of parent_state.
        """
        new_solution_prefix = (parent_state.solution_prefix + '\n\n' + rollout).strip() if parent_state.solution_prefix else rollout
        new_state = State(solution_prefix=new_solution_prefix, parent=parent_state)
        new_state.MC = 1.0  # Since the rollout is correct
        new_state.total_rollouts = 0
        new_state.correct_rollouts = 0
        self.T.add_state(new_state)
        parent_state.children.append(new_state)  # Add to parent's children

    def add_incorrect_rollout_to_tree(self, parent_state: State, rollout: str):
        """
        Add the incorrect rollout to the tree as a child of parent_state.

        Parameters:
        - parent_state (State): The state from which the rollout was selected.
        - rollout (str): The incorrect rollout string.
        """
        new_solution_prefix = (parent_state.solution_prefix + '\n\n' + rollout).strip() if parent_state.solution_prefix else rollout
        new_state = State(solution_prefix=new_solution_prefix, parent=parent_state)
        new_state.MC = 0.0  # Since the rollout is incorrect
        new_state.total_rollouts = 0
        new_state.correct_rollouts = 0
        self.T.add_state(new_state)
        parent_state.children.append(new_state)  # Add to parent's children

    async def binary_search_incorrect_step(self, s_ast: State, steps: List[str], left: int, right: int):
        """
        Recursively perform binary search to find all incorrect steps in the rollout.
        """
        if left > right:
            return

        mid = (left + right) // 2
        new_steps = steps[left:mid + 1]
        if new_steps:
            prefix_solution = s_ast.solution_prefix + '\n\n' + separate_steps(new_steps, mode='join')
        else:
            prefix_solution = s_ast.solution_prefix
        # Create new state s_new
        s_new = State(solution_prefix=prefix_solution.strip(), parent=s_ast)
        self.T.add_state(s_new)
        s_ast.children.append(s_new)

        # Perform Monte Carlo estimation for s_new
        await self.monte_carlo_estimation(s_new)

        if s_new.MC == 0:
            # Found incorrect step; continue searching in the left half to find earlier incorrect steps
            await self.binary_search_incorrect_step(s_ast, steps, left, mid - 1)
        else:
            # Steps up to mid are correct; continue searching in the right half
            await self.binary_search_incorrect_step(s_new, steps, mid + 1, right)

    async def expansion_phase_binary_search(self, parent_state: State, rollout: str):
        """
        Expansion phase that adds the rollout as a new state and performs Monte Carlo estimation
        using Binary Search to efficiently find the correct rollout.
        """
        # Separate the rollout into individual steps
        steps = separate_steps(rollout, mode='split')

        # Perform binary search to find incorrect steps
        await self.binary_search_incorrect_step(parent_state, steps, 0, len(steps) - 1)

    def maintenance_phase(self, state: State):
        """
        Update statistics and candidate pool for all incorrect rollouts associated with the state.

        Parameters:
        - state (State): The state whose incorrect rollouts need to be updated.
        """

        # Iterate through all incorrect rollouts of the state
        for rollout in state.incorrect_rollouts:
            # Since we've already determined these rollouts are incorrect, no need to re-evaluate correctness

            priority = self.compute_selection_score(state, rollout)
            # Update the candidate pool with the new priority
            self.C.add_or_update(state, rollout, priority)
            # print(f"Updated Incorrect Rollout: '{rollout}' with new priority: {priority:.4f}")

        # print("Maintenance Phase Completed.\n")

    def collect_solution_prefixes(self) -> List[Dict[str, Any]]:
        """
        Collect all solution prefixes and their corresponding MC values from the search tree.

        Returns:
            List[Dict[str, Any]]: A list of dictionaries containing solution prefixes and their MC values.
        """
        collected_data = []
        for node in self.T.nodes:
            solution_prefix = node.solution_prefix
            mc_value = node.MC
            collected_data.append({
                "solution_prefix": solution_prefix,
                "mc_value": mc_value
            })
        return collected_data

    def collect_tree_structure(self) -> Dict[str, Any]:
        """
        Collect the tree structure starting from the root.

        Returns:
            Dict[str, Any]: A nested dictionary representing the tree structure.
        """
        if self.T.root:
            tree_data = self.T.root.get_text_with_labels()
            return tree_data
        return {}


# Example usage
if __name__ == "__main__":
    # Initialize the Language Model's AsyncOpenAI client for LM.
    from openai import AsyncOpenAI
    lm_client = AsyncOpenAI(
        base_url="http://localhost:8000/v1",
        api_key="token-abc123",
    )
    
    LM = LanguageModel(
        client=lm_client,
        max_new_tokens=4096,
        temperature=0.7,
        top_p=0.9,
        model_name="DeepSeek-R1-Distill-Qwen-14B"
    )

    # Define the question and expected answer
    question = "Melinda will roll two standard six-sided dice and make a two-digit number with the two numbers she rolls. For example, if she rolls a 6 and a 3, she can either form 36 or 63. What is the probability that she will be able to make an integer between 10 and 20, inclusive? Express your answer as a common fraction."
    expected_answer =  "\\frac{11}{36}"

    client = AsyncOpenAI(
        base_url="http://localhost:8001/v1",
        api_key="token-abc123",
    )  # This is a placeholder; ensure client supports sync chat.completions.create
    reward_model = ProcessRewardModel(client, model="deepseek-14b-prm-filtered-balance-full", temperature=0.0, max_tokens=1)

    # Initialize OmegaPRM with parameters and the reward model instance
    omega_prm = OmegaPRM(
        LM=LM,
        reward_model=reward_model,
        c_puct=0.125,
        alpha=0.5,
        beta=0.9,
        L=500,
        k=8,
        N=10,
        rollout_budget=20,
        save_data_tree=True,
    )

    # Run the OmegaPRM algorithm
    collected_data = asyncio.run(omega_prm.run(question, expected_answer))

    # Save the collected solutions to a JSON file
    with open("collected_solutions2.json", "w") as f:
        json.dump(collected_data, f, indent=4)