From ee87ca852434ae0a8dd88009fa997ca64279fae3 Mon Sep 17 00:00:00 2001 From: YanSong97 Date: Tue, 5 Nov 2024 23:42:15 +0000 Subject: [PATCH] add critic_mcts; add full_answer property to base_env; p.evaluate_problem also output metadata; --- envs/MATH/parse_utils_qwen.py | 77 ++++-- envs/base_env.py | 7 + envs/critic_MATH/__init__.py | 3 + envs/critic_MATH/critic_math.py | 390 ++++++++++++++++++++++++++++ envs/critic_MATH/data.py | 26 ++ envs/critic_MATH/prompt.py | 14 + reason/evaluation/evaluate.py | 12 +- reason/evaluation/evaluator.py | 16 +- reason/evaluation/methods.py | 70 +++++ reason/guided_search/critic_mcts.py | 221 ++++++++++++++++ reason/guided_search/tree.py | 10 +- scripts/eval/critic_mcts.sh | 14 + 12 files changed, 830 insertions(+), 30 deletions(-) create mode 100644 envs/critic_MATH/__init__.py create mode 100644 envs/critic_MATH/critic_math.py create mode 100644 envs/critic_MATH/data.py create mode 100644 envs/critic_MATH/prompt.py create mode 100644 reason/guided_search/critic_mcts.py create mode 100644 scripts/eval/critic_mcts.sh diff --git a/envs/MATH/parse_utils_qwen.py b/envs/MATH/parse_utils_qwen.py index e65a943..4d23d0c 100644 --- a/envs/MATH/parse_utils_qwen.py +++ b/envs/MATH/parse_utils_qwen.py @@ -482,6 +482,32 @@ def extract_theoremqa_answer(pred: str, answer_flag: bool = True): return pred +def check_boxed(pred_str): + ans = pred_str.split("boxed")[-1] + if len(ans) == 0: + return "" + elif ans[0] == "{": + stack = 1 + a = "" + for c in ans[1:]: + if c == "{": + stack += 1 + a += c + elif c == "}": + stack -= 1 + if stack == 0: + break + a += c + else: + a += c + else: + a = ans.split("$")[0].strip() + + # if the answer is a equation + if "=" in a: + a = a.split('=')[-1] + + return a def extract_answer(pred_str, data_name, use_last_number=True): pred_str = pred_str.replace("\u043a\u0438", "") @@ -490,27 +516,16 @@ def extract_answer(pred_str, data_name, use_last_number=True): # minerva_math tmp = pred_str.split("final answer is $", 1)[1] pred = tmp.split("$. I hope", 1)[0].strip() + if "boxed" in pred: # llama3 case, check boxed + pred = check_boxed(pred) + elif "final answer is" in pred_str and ". I hope" in pred_str: # llama3-8b-instruct case + tmp = pred_str.split("final answer is", 1)[1] + pred = tmp.split(". I hope", 1)[0].strip() + elif "final answer is" in pred_str: + tmp = pred_str.split("final answer is", 1)[1] + pred = tmp.strip() elif "boxed" in pred_str: - ans = pred_str.split("boxed")[-1] - if len(ans) == 0: - return "" - elif ans[0] == "{": - stack = 1 - a = "" - for c in ans[1:]: - if c == "{": - stack += 1 - a += c - elif c == "}": - stack -= 1 - if stack == 0: - break - a += c - else: - a += c - else: - a = ans.split("$")[0].strip() - pred = a + pred = check_boxed(pred_str) elif "he answer is" in pred_str: pred = pred_str.split("he answer is")[-1].strip() elif "final answer is" in pred_str: @@ -617,3 +632,25 @@ def parse_question(example, data_name): question += " (Yes or No)" return question.strip() +def extract_groundtruth(groundtruth_str: str) -> str: + return parse_ground_truth(groundtruth_str, data_name='math') + + +if __name__ == "__main__": + # run examples + # test_text = "n\nLet's figure out how much the math club made:\n\\begin{align*} \n &\\text{Cookies: } \\frac{54}{3}=18 \\text{ sets of } 3 cookies, \\text{ so } 18 \\cdot \\$1 = \\$18. \\\\\n &\\text{Cupcakes: } 20 \\cdot \\$2 = \\$40. \\\\\n &\\text{Brownies: } 35 \\cdot \\$1 = \\$35. \\\\\n &\\text{Total: } \\$18 + \\$40 + \\$35 = \\$93.\n\\end{align*}The math club spent $\\$15$to bake these items, so their profit is $\\boxed{\\$93-\\$15 = \\$78}.$\n\n" + # test_text = "\nStep 1: List all the possible ways to insert parentheses. There are 6 possible ways to insert 3 pairs of parentheses into 4 terms. The 6 ways are: \\begin{align*}\n&(2\\cdot 3)\\cdot(4\\cdot 5)+1\\\\\n&2\\cdot(3\\cdot 4)\\cdot (5+1)\\\\\n&(2\\cdot 3)\\cdot(4\\cdot (5+1))\\\\\n&(2\\cdot (3\\cdot 4)\\cdot 5)+1\\\\\n&2\\cdot (3\\cdot (4\\cdot 5))+1\\\\\n&(2\\cdot ((3\\cdot 4)\\cdot 5))+1\n\\end{align*} Step \"\nFinal Answer: The final answer is 6. I hope it is correct.\n" + # test_text = "\nWe can simplify $f(x)$ as follows:\n$f(x) = \\frac{3x-2}{x-2} = \\frac{3x-2}{x-2} \\cdot \\frac{x+2}{x+2} = \\frac{3x^2 - 2x - 6}{x^2-4} = 3 - \\frac{8}{x^2-4}$.\nThen, we can find the values of $f(-2)$, $f(-1)$, and $f(0)$:\n$f(-2) = 3 - \\frac{8}{(-2)^2-4} = 3 - \\frac{8}{4} = 1$,\n$f(-1) = 3 - \\frac{8}{(-1)^2-4} = 3 - \\frac{8}{1-4} = 3 - (-8) = 11$,\n$f(0) = 3 - \\frac{8}{0^2-4}$ is undefined, because the denominator is zero.\nSo, the sum of the values is: $\\boxed{\\frac{1 + 11}{1} = 12}$.\n" + # test_text = "\nLet $n = 7k+2$ for some integer $k$. Then we have\n\\begin{align*}\n(n+2)(n+4)(n+6) & = (7k+2+2)(7k+2+4)(7k+2+6) \\\\\n& = (7k+4)(7k+6)(7k+8)\n\\end{align*}By Fermat's Little Theorem, $7^3 \\equiv 1 \\pmod{7}$. Therefore, we have\n$$(7k+4)(7k+6)(7k+8) = 7k(7k+1)(7k+2)$$ $$\\equiv k(k+1)(k+2) \\pmod{7}$$ $$\\equiv k(k+1)(k+3) \\pmod{7}$$ $$\\equiv k(k+1)(-1) \\pmod{7}$$ $$\\equiv -k(k+1) \\pmod{7}$$ $$\\equiv -(k^2+k) \\pmod{7}$$ $$\\equiv -(k^2+k-1+1) \\pmod{7}$$ $$\\equiv -(k^2+k-1)+1 \\pmod{7}$$ $$\\equiv -(k^2+k) + (k-1) + 1 \\pmod{7}$$ $$\\equiv (-k)(k-1) + (k-1) + 1 \\pmod{7}$$ $$\\equiv -k^2+k+k-1+1 \\pmod{7}$$ $$\\equiv -k^2+2k \\pmod{7}$$ $$\\equiv -k^2+2k-2+2 \\pmod{7}$$ $$\\equiv -(k^2+2k-2) + 2 \\pmod{7}$$ $$\\equiv -(k+1)^2+1+2 \\pmod{7}$$ $$\\equiv -(k+1)^2+3 \\pmod{7}$$ $$\\equiv -(k+1)(k+1)+3 \\pmod{7}$$ $$\\equiv -(k+1)(k+1)+2+1 \\pmod{7}$$ $$\\equiv -(k+1)(k+1)+(k+1)+1 \\pmod{7}$$ $$\\equiv -((k+1)^2+(k+1)) + 1 \\pmod{7}$$ $$\\equiv -((k+1)(k+1)+k+1) + 1 \\pmod{7}$$ $$\\equiv -((k+1)(k+1)+k+1-1+1) \\pmod{7}$$ $$\\equiv -((k+1)(k+1)+k) + 2 \\pmod{7}$$ $$\\equiv -((k+1)(k+1)-k)+2 \\pmod{7}$$ $$\\equiv -(k+1)^2+2 \\pmod{7}$$ $$\\equiv -(k+1)^2+2-2+2 \\pmod{7}$$ $$\\equiv -(k+1)^2+0+2 \\pmod{7}$$ $$\\equiv -(k+1)^2+2 \\pmod{7}$$ $$\\equiv -1^2+2 \\pmod{7}$$ $$\\equiv -1+2 \\pmod{7}$$ $$\\equiv 1 \\pmod{7}$$ $$\\boxed{\\equiv 1}$$\nFinal Answer: The final answer is 1. I hope it is correct.\n" + # test_text = "\\boxed{\\text{Step 1:} } x^6 - 3 = x^6 + 0x^5 + 0x^4 + 0x^3 + 0x^2 + 0x - 3. \n\\boxed{\\text{Step 2:} } The dividend has been written in descending powers of $x$. Let us long divide.\n\\boxed{\\text{Step 3:} } \\begin{array}{r} x^5- x^4+x^3-x^2+x-1 \\\\ x + 1 \\enclose{longdiv}{ x^6-0x^5+0x^4+0x^3+0x^2-3} \\\\ \\underline{x^6+ x^5} \\\\ -x^5+0x^4+0x^3+0x^2-3 \\\\ \\underline{-x^5- x^4} \\\\ x^4+0x^3+0x^2-3 \\\\ \\underline{x^4+ x^3} \\\\ -x^3+0x^2-3 \\\\ \\underline{-x^3- x^2} \\\\ x^2-3 \\\\ \\underline{x^2+ x} \\\\ -x-3 \\\\ \\underline{-x-1} \\\\ -2 \\end{array} \n\\boxed{\\text{Step 4:} } The quotient is $x^5- x^4+x^3-x^2+x-1$. The remainder is $-2$. \n\\boxed{\\text{Final Answer: } x^5- x^4+x^3-x^2+x-1}\n" + # test_text = "\\step{Step 1}\nWe are given that $2^8=4^x$. To solve for $x$, we can rewrite $4$ as $2^2$. \n$2^8=(2^2)^x$\nWe can now rewrite the equation as:\n$\\boxed{(2^2)^x=2^8}$\n\\step{Step 2}\nWe know that the exponent of $2$ is equal on both sides of the equation, so we can equate them:\n$2^x=2^8$\nTo solve for $x$, we can take the logarithm base 2 of both sides:\n$\\log_22^x=\\log_22^8$\nUsing the property of logarithms that $\\log_ab^x=x$, we get:\n$x\\log_22=8\\log_22$\nSimplifying:\n$x=8$\n\\step{Final Answer}\nThe final answer is $\\boxed{x=8}$. I hope it is correct.\n" + # test_text = "\n\\begin{step} \"Step 1: Split the equation into two cases\"\nThe equation $|5x-1|=x+3$ can be written as either $5x-1=x+3$ or $-(5x-1)=x+3$. The first equation becomes $4x=4$, which has the solution $x=1$. The second equation becomes $-5x+1=x+3$, which has the solution $x=-2$. So we have two solutions, $x=1$ and $x=-2$.\n\\end{step}\n\\begin{step} \"Step 2: Determine the larger solution\"\nThe solution $x=1$ is greater than $x=-2$, so it is the larger solution.\n\\end{step}\n\\begin{boxed} $\\boxed{x=1}$ \\end{boxed}\nFinal Answer: The final answer is 1. I hope it is correct.\n" + # test_text = "\\boxed{\\text{Step 1:}}\nThere are $6\\times 6=36$ possible outcomes for the two dice. The product of the numbers on the dice can only be a multiple of 5 if one of the numbers is 5 and the other number is 1, or if both numbers are 5. There are $2\\times1=2$ ways to get the first outcome, and one way to get the second outcome.\n\\boxed{\\text{Step 2:}}Therefore, there are 3 ways for the product to be a multiple of 5. The probability of the product being a multiple of 5 is then \\begin{align*}\n\\frac{3}{36}&=\\frac{1}{12}.\n\\end{align*}\nFinal Answer: The final answer is \\frac{1}{12}.\n" + test_text = "}\\boxed{\\text{Step 1: Count the number of ways to choose 3 boys}}\\text{Mr. Brennan has 7 boys and he needs to choose 3 of them. We can do this in }\\binom{7}{3}=\\frac{7!}{3!(7-3)!}=\\frac{7!}{3!4!}=\\frac{7\\cdot6\\cdot5}{3\\cdot2\\cdot1}=35\\text{ ways.}\\text{Step }\"\n\\}\\boxed{\\text{Step 2: Count the number of ways to choose 2 girls}}\\text{Mr. Brennan has 4 girls and he needs to choose 2 of them. We can do this in }\\binom{4}{2}=\\frac{4!}{2!(4-2)!}=\\frac{4!}{2!2!}=\\frac{4\\cdot3}{2\\cdot1}=6\\text{ ways.}\\text{Step }\"\n\\}\\boxed{\\text{Step 3: Combine the results}}\\text{The total number of ways Mr. Brennan can pick 3 boys and 2 girls is the product of the number of ways to choose each set. Therefore, the total number of ways is }35\\cdot 6=210.\\text{Step }\"\nFinal Answer: The final answer is 210. I hope it is correct.\n" + + # true_text = "To find the profit, we want to find out how much the math club earned from selling the various baked goods and subtract the cost of producing those goods, $\\$15$, from the number we get.\n\nFirst let's calculate how much the math club earned from selling cookies. The cookies were sold at a price of three for $\\$1$, so the math club earned $54\\div 3\\cdot\\$1=18\\cdot\\$1=\\$18$ from selling cookies.\n\nNext, let's calculate how much the club earned from selling cupcakes. At a price of $\\$2$ each, the club earned $20\\cdot \\$2=\\$40$ from selling cupcakes.\n\nFinally, let's calculate how much the club earned from selling brownies. At a price of $\\$1$ each, the club earned $35\\cdot\\$1=\\$35$ from selling brownies.\n\nNow let's add up these numbers to find out how much the club earned in total and subtract $\\$15$ from that number to find the club's profit. We obtain \\begin{align*}\n\\$18+\\$40+\\$35-\\$15&=\\$18+\\$40+\\$35-\\$15\\\\\n&=\\$18+\\$40+\\$35+(-\\$15)\\\\\n&=\\$18+\\$40+(\\$35+(-\\$15))\\\\\n&=\\$18+\\$40+(\\$20)\\\\\n&=\\boxed{78}.\n\\end{align*}Notice how we used the definition of subtraction, $a-b=a+(-b)$ to $\\$35-\\$15$ as $\\$35+(-\\$15)$ and the associative property of addition to group the numbers together." + true_text = "\nStep 1: The skater spins 2250 degrees to her right, which means she turns 6.25 full rotations to her right (since 2250 degrees is equal to 6.25 * 360 degrees).\nStep 2: Since she starts facing north, after each full rotation, she will be facing the original direction (north) again. So, after 6.25 full rotations, she will still be facing the original direction, which is north, but rotated 6.25 * 90 = 562.5 degrees to the right.\nStep 3: Since 562.5 degrees is equivalent to 1.5625 full rotations, and she turns 1.5625 * 90 = 140.625 degrees to the right, she will be facing slightly east of north.\nStep 4: So, when she finishes her spin, she will be facing east, since east is slightly more than 90 degrees to the right of north.\n\\boxed{East}\n" + + answer = extract_answer(test_text, "MATH") + true_answer = extract_groundtruth(true_text) + print(f"answer = {answer}, true label = {true_answer}") \ No newline at end of file diff --git a/envs/base_env.py b/envs/base_env.py index 3047cc9..03337e7 100644 --- a/envs/base_env.py +++ b/envs/base_env.py @@ -266,6 +266,13 @@ def question(self)->str: @property def answer(self): + """ + partial answer + """ + return "".join(self.action_history) + + @property + def full_answer(self): return "".join(self.action_history) def get_done_and_info(self): diff --git a/envs/critic_MATH/__init__.py b/envs/critic_MATH/__init__.py new file mode 100644 index 0000000..0a117fb --- /dev/null +++ b/envs/critic_MATH/__init__.py @@ -0,0 +1,3 @@ +from .critic_math import Env +from envs.MATH.env import extract_answer, extract_groundtruth, judge_correct +from .data import get_train_test_dataset diff --git a/envs/critic_MATH/critic_math.py b/envs/critic_MATH/critic_math.py new file mode 100644 index 0000000..bf6d597 --- /dev/null +++ b/envs/critic_MATH/critic_math.py @@ -0,0 +1,390 @@ +from typing import List, Dict, Tuple +import os +import numpy as np +import json + +from envs.MATH.env import CoTEnv +from envs.base_env import NoLegalActionException, ResetException +from tqdm import tqdm +from envs.MATH.env import extract_answer, extract_groundtruth, judge_correct +from reason.inference.lm_call import LMCallingConfig +from .prompt import ( + COT_TASK_DESC, + CRITIQUE_TASK_DESC, + REWRITE_TASK_DESC, + COT_FORMAT_STR, + CRITIQUE_FORMAT_STR, + REWRITE_FORMAT_STR, + SEP, +) +from distributed.utils import print_with_rank +from loguru import logger + +from pathlib import Path + +# Get the file path of the current script +CURRENT_DIR = Path(__file__).parent + + +ANS_RE = None +STOP_STR = None + + +def read_txt(file_path): + assert str(file_path).endswith(".txt") + with open(file_path, "r", encoding="utf-8") as f: + data = f.read() + return data + + +def read_json(file_path): + assert str(file_path).endswith(".json") + with open(file_path, "r", encoding="utf-8") as f: + data = json.load(f) + return data + + +class Env(CoTEnv): + def __init__( + self, + config, + math_problems, + llm_gen_fn, + task_desc_str: str = None, + cot_example_str: str = None, + problem_format_str: str = None, + cot_task_desc_str: str = COT_TASK_DESC, + critique_task_desc_str: str = CRITIQUE_TASK_DESC, + rewrite_task_desc_str: str = REWRITE_TASK_DESC, + cot_format_str: str = COT_FORMAT_STR, + critique_format_str: str = CRITIQUE_FORMAT_STR, + rewrite_format_str: str = REWRITE_FORMAT_STR, + reset=False, + ): + """ + three types of thinking in reasoning: + 1. i.i.d sampling: generative independent response + 2. conditional sampling (counterfactual): what if the response is wrong? + 3. reflective sampling + Args: + config: + math_problems: + llm_gen_fn: + task_desc_str: + cot_example_str: + problem_format_str: + reset: + """ + super().__init__( + config, + math_problems, + llm_gen_fn, + task_desc_str, + cot_example_str, + problem_format_str, + reset=False, + ) + + self.current_node_type = None + # LLM generation config + self.gen_cfg = config["generation_config"] + + # default parameter setting + self.num_first_answer = 3 + self.num_reviews = 3 + self.num_rewrite = 3 + self.max_new_tokens_answer = 1024 + self.max_new_tokens_review = 1024 + self.max_depth = 5 + + self.print_log = True + self.total_api_call_completion = 0 + self.total_tree_completion = 0 + + # self.task_name = "MATH" + self.sep = None + self._init_query = None + self._next_state_terminated = None + + # loading template + # with open(os.path.join( + # CURRENT_DIR, f"prompts/MATH/decompose/decompose_template.json"), + # "r") as f: + # decompose_template = json.load(f) + # self.question_index = decompose_template["index"] + # self.critic_prompt_template = read_json( + # os.path.join(CURRENT_DIR, f"prompts/MATH/generation_config.json")) + self.cot_task_desc = cot_task_desc_str + self.critique_task_desc = critique_task_desc_str + self.rewrite_task_desc = rewrite_task_desc_str + + self.cot_format_str = cot_format_str + self.critique_format_str = critique_format_str + self.rewrite_format_str = rewrite_format_str + + if reset: + self.reset(update_legal_action=True) + # self.rewrite_prompt_template = None + + def check_attribute(self): + assert hasattr(self, "cot_task_desc") + assert hasattr(self, "critique_task_desc") + assert hasattr(self, "rewrite_task_desc") + assert hasattr(self, "cot_format_str") + assert hasattr(self, "critique_format_str") + assert hasattr(self, "rewrite_format_str") + + @property + def stop_str(self): + return STOP_STR + + @property + def answer(self): + return "" + + @property + def full_answer(self): + return self.action_history[-1] + + def post_process_act(self, action: str): + if not action.endswith(self.sep): + action = action.strip() + self.sep + + return action + + def set_problem(self, idx): + self.math_problem = self.math_problems[idx] + + def get_state(self): + # if no solution has been generated yet, generate the initial query + if len(self.action_history) == 0: + ret = self.cot_task_desc + self.cot_format_str.format( + question=self.question + ) + else: + ret = ( + self.cot_task_desc + + self.cot_format_str.format(question=self.question) + + self.action_history[-1] + ) + return ret + + def reset(self, update_legal_action=True): + """ + reset the environment, and generate the first solution to the question + Args: + update_legal_action: + + Returns: + + """ + assert update_legal_action, print("Need to first update legal action") + self.set_problem(idx=0) + self.action_history = [] + self.review_history = [] + + if update_legal_action: + cnt = 0 + while cnt < 3: + cnt += 1 + try: + self._legal_actions, api_completion_token = ( + self.generate_first_response() + ) + break + except NoLegalActionException as e: + if cnt == 3: + raise ResetException + info = {"api_completion_token": api_completion_token} + return None, info + + def generate_first_response(self) -> (List[Dict], int): + first_cot_prompt = self.cot_task_desc + self.cot_format_str.format( + question=self.math_problem["question"] + ) + result = self.llm_gen_fn( + input_str=first_cot_prompt, + config=LMCallingConfig( + n=self.num_first_answer, + stop_str=self.sep, + include_stop_str_in_output=True, + max_new_tokens=self.max_new_tokens_answer, + **self.config["generation_config"] + ), + ) + + texts = result.text + logps_avg_by_len = result.logp_avg_by_len + token_len = result.num_tokens + + # for i in texts: + # print(f"\n {i}") + + _legal_actions = [ + { + "action": action, + "prob": prob, + "num_token": n_token, + "finish_reason": finish_reason, + "from_review": "", + } + for action, prob, n_token, finish_reason in zip( + texts, logps_avg_by_len, token_len, result.finish_reason + ) + ] + self._next_state_terminated = dict(zip(texts, [False] * len(texts))) + return _legal_actions, result.completion_tokens + + def step(self, action, update_legal_action=True): + """ + + Args: + action: the chosen action, which is the refined solution in this case, need to record this + update_legal_action: + + Returns: + + """ + self.action_history.append(action) # recording all the select full answer + + reward = self.get_reward() + terminated, truncated, info = ( + self.get_done_and_info() + ) # terminated or truncated when reach maximim depth + # update legal actions + if not (terminated or truncated) and update_legal_action: + cnt = 0 + while cnt < 3: + cnt += 1 + try: + self._legal_actions, api_completion_token = ( + self.update_legal_actions() + ) + info["api_completion_token"] = api_completion_token + except NoLegalActionException as e: + if cnt == 3: + terminated = True + reward = 0 + self._legal_actions = None + info["winner"] = 2 + info["api_completion_token"] = 0 + else: + pass + else: + self._legal_actions = None + if info["winner"] == 1: + reward = 1.0 + info["api_completion_token"] = 0 + return None, reward, terminated, truncated, info + + def update_legal_actions(self): + """ + Given the current state (the completed solution to Q), a critic LLM generate review and + a proposer LLM rewrite the answer, which is the new updated legal action + Returns: + + """ + # retrive current answer + assert len(self.action_history) > 0 + current_action = self.action_history[-1] + + # review + review_prompt = self.critique_task_desc + self.critique_format_str.format( + question=self.math_problem["question"], answer=current_action + ) + result = self.llm_gen_fn( + input_str=review_prompt, + config=LMCallingConfig( + n=self.num_reviews, + stop_str=self.sep, + include_stop_str_in_output=True, + max_new_tokens=self.max_new_tokens_review, + **self.config["generation_config"] + ), + ) + + review_texts = result.text + # for i in review_texts: + # print(f"\n review = {i}") + # logps_avg_by_len = result.logp_avg_by_len + # token_len = result.num_tokens + + new_action_text = [] + from_review_text = [] + new_prob_list = [] + tokens_num_list = [] + final_reason_list = [] + total_completion = result.completion_tokens + for text in review_texts: + rewrite_prompt = self.rewrite_task_desc + self.rewrite_format_str.format( + question=self.math_problem["question"], + answer=current_action, + review=text, + ) + + result = self.llm_gen_fn( + input_str=rewrite_prompt, + config=LMCallingConfig( + n=self.num_rewrite, + stop_str=self.sep, + include_stop_str_in_output=True, + max_new_tokens=self.max_new_tokens_answer, + **self.config["generation_config"] + ), + ) + + new_action_text += result.text + new_prob_list += result.logp_avg_by_len + tokens_num_list += result.num_tokens + final_reason_list += result.finish_reason + total_completion += result.completion_tokens + from_review_text.append(text) + + _legal_actions = [ + { + "action": action, + "prob": prob, + "num_token": n_token, + "finish_reason": finish_reason, + "from_review": f_review, + } + for action, prob, n_token, finish_reason, f_review in zip( + new_action_text, + new_prob_list, + tokens_num_list, + final_reason_list, + from_review_text, + ) + ] + + self._next_state_terminated = dict( + zip(new_action_text, [False] * len(new_action_text)) + ) + + return _legal_actions, total_completion + + def get_done_and_info(self): + info = {"winner": 0} + # done when reaches maximum length + truncated = terminated = len(self.action_history) >= self.config["max_length"] + assert len(self.action_history) <= self.config["max_length"] + if terminated or truncated: + if self._is_correct(self.action_history[-1]): + info["winner"] = 1 + else: + info["winner"] = 2 + return terminated, truncated, info + return terminated, truncated, info + + def _is_correct(self, completion): + extracted_answer = extract_answer(completion) + # print("Compare: {} -- {}".format(extrated_answer, + # self.math_problem['answer'])) + # return extrated_answer == self.math_problem['answer'] + return judge_correct( + self.math_problem["question"], self.math_problem["answer"], extracted_answer + ) + + def get_reward(self): + """To implement based on learned reward model""" + return 0 diff --git a/envs/critic_MATH/data.py b/envs/critic_MATH/data.py new file mode 100644 index 0000000..165cdcf --- /dev/null +++ b/envs/critic_MATH/data.py @@ -0,0 +1,26 @@ +from pathlib import Path +import jsonlines +from torch.utils.data import Dataset + + +def get_train_test_dataset(*args, **kwargs): + env_dir = Path(__file__).parent.parent + test_ds = JsonlMathDataset(env_dir / "MATH/dataset/test500.jsonl") + train_ds = JsonlMathDataset(env_dir / "MATH/dataset/train.jsonl") + return train_ds, test_ds + + +class JsonlMathDataset(Dataset): + def __init__(self, data_path): + super().__init__() + self.data = [] + with jsonlines.open(data_path, "r") as reader: + for obj in reader: + self.data.append(obj) + + def __len__(self): + return len(self.data) + + def __getitem__(self, index): + x = self.data[index] + return {"question": x["problem"], "answer": x["solution"]} diff --git a/envs/critic_MATH/prompt.py b/envs/critic_MATH/prompt.py new file mode 100644 index 0000000..d56ec7e --- /dev/null +++ b/envs/critic_MATH/prompt.py @@ -0,0 +1,14 @@ +COT_TASK_DESC = '<|im_start|>system\nPlease wrie a answer for this question. The Answer should format as "\n#### Reasoning Process\n...\n#### Verification\n...\n#### Final Answer\n...\n"The Final Answer should format as \\boxed{{}}.<|im_end|>' + +CRITIQUE_TASK_DESC = "<|im_start|>system\nPlease analyze this weak Answer, write a strict Critic/Reflection for error re-correct and Hints/Guidelines for maximum improvement. Let’s think step by step.<|im_end|>" + +REWRITE_TASK_DESC = "<|im_start|>system\nPlease write a better answer for this question refer to the comments. Please reason step by step, and put your final answer within \\boxed{{}}.<|im_end|>" + +COT_FORMAT_STR = ( + """<|im_start|>user\nQuestion: {question}<|im_end|>\n<|im_start|>assistant\n""" +) +CRITIQUE_FORMAT_STR = """<|im_start|>user\nQuestion: {question}\nAnswer: {answer}<|im_end|>\n<|im_start|>assistant\n""" +REWRITE_FORMAT_STR = """<|im_start|>user\nQuestion: {question}\nAnswer: {answer}\nReview: {review}<|im_end|>\n<|im_start|>assistant\n""" + + +SEP = None diff --git a/reason/evaluation/evaluate.py b/reason/evaluation/evaluate.py index 17cf6b1..01d8d7f 100644 --- a/reason/evaluation/evaluate.py +++ b/reason/evaluation/evaluate.py @@ -150,7 +150,7 @@ def parallel_evaluate_test_dataset( ) # Distributes tasks from the test_ds dataset across the worker pool asynchronously and # collects results in any order as they complete. Every worker has a new searching tree as we reset the # tree in solver_fn - for i, (problem_inst, result, output) in enumerate( + for i, (problem_inst, result, output, metadata_dict) in enumerate( tqdm(res_q, total=len(test_ds)) ): results.append(result) @@ -161,6 +161,7 @@ def parallel_evaluate_test_dataset( "groundtruth": problem_inst["answer"], "result": result, "output": output, + **metadata_dict } record_writer.write(obj) avg_res = (tree.map_structure(lambda *xs: np.mean(xs), *results),) @@ -218,6 +219,15 @@ def parallel_evaluate_test_dataset( num_path=config.num_sequence, ) solver_fn = partial(rstar_mcts, method_config, gen_config) + elif config.method == "critic_mcts": + method_config = VanilaMCTSConfig( + task_name=config.task_name, + tree_max_depth=config.tree_max_depth, + tree_max_width=config.tree_max_width, + select_by_prior=False, + num_path=config.num_sequence, + ) + solver_fn = partial(critic_mcts, method_config, gen_config) else: raise ValueError(f"Unknown method: {config.method}") diff --git a/reason/evaluation/evaluator.py b/reason/evaluation/evaluator.py index 48111e1..5fd225f 100644 --- a/reason/evaluation/evaluator.py +++ b/reason/evaluation/evaluator.py @@ -94,12 +94,12 @@ class SolutionOutput: # for beam search, it's a list of zeros, except the last element indicates total tokens # for mcts, it's a list of int, indicate how many tokens comsumed between two paths completion_tokens: List[int] + metadata: Optional[Dict] = None # other information @dataclass class TreeSearchSolutionOutput(SolutionOutput): - tree_completion_tokens: List[int] - + tree_completion_tokens: List[int]=None class MathEvaluator: @@ -119,7 +119,7 @@ def __init__( def evaluate_problem( self, problem_inst: Dict[str, str], solver_fn: Callable - ) -> List[str]: + ) -> (str, str, str, Dict[str, str]): solution: SolutionOutput = solver_fn(problem_inst, self.lm_call, self.rm_call) result, output = self.analyze_output(problem_inst, solution.solutions) total_completion_token = 0 @@ -131,7 +131,9 @@ def evaluate_problem( # answers, therefore we need to take sum here. total_completion_token += solution.completion_tokens[i] result["total_completion_tokens"] = total_completion_token - return problem_inst, result, output + + processed_metadata = self.process_meta_data(solution.metadata) + return problem_inst, result, output, processed_metadata def analyze_output(self, problem_inst: Dict[str, str], gen_answers: List[str]): extracted_groundtruth = self._task.extract_groundtruth(problem_inst["answer"]) @@ -163,6 +165,12 @@ def analyze_output(self, problem_inst: Dict[str, str], gen_answers: List[str]): } return res, output_list + def process_meta_data(self, metadata_dict=None): + if metadata_dict is None: + return {} + + return metadata_dict + @ray.remote class RemoteMathEvaluator(MathEvaluator): diff --git a/reason/evaluation/methods.py b/reason/evaluation/methods.py index bad3bd7..bc173d6 100644 --- a/reason/evaluation/methods.py +++ b/reason/evaluation/methods.py @@ -6,6 +6,7 @@ from reason.evaluation.evaluator import SolutionOutput, Task, TreeSearchSolutionOutput from reason.guided_search.tree import SearchTree from reason.guided_search.rstar import RstarSearchTree +from reason.guided_search.critic_mcts import CriticSearchTree @dataclass @@ -261,3 +262,72 @@ def rstar_mcts( completion_tokens=[t["api_completion_tokens"] for t in traj_list], tree_completion_tokens=[t["tree_completion_tokens"] for t in traj_list], ) + + +@dataclass +class CriticMCTSConfig(MCTSBaseConfig): + # rollout step strategy, if `select_by_prior` is False, + # then select by the initial critic value + # otherwise, random choice by the prior probability + select_by_prior: bool = False + num_path: int = 1 + + def __post_init__(self): + super().__post_init__() + if not self.select_by_prior: + assert self.init_critic_value, \ + "VanilaMCTS with greedy as rollout method should set init_critic_value to True" + assert self.num_path > 0 + + +def critic_mcts( + config: CriticMCTSConfig, + gen_config: LMCallingConfig, + problem_inst: Dict[str, str], + lm_call: LanguageModelCallingFunction, + rm_call: RewardModelCallingFunction +): + task = Task(task_name=config.task_name) + env = task.env_fn( + config={ + "max_actions": config.tree_max_width, + "max_length": config.tree_max_depth, + "stop_str": "The answer is ", + "generation_config": { + "temperature": gen_config.temperature, + "top_p": gen_config.top_p, + "top_k": gen_config.top_k, # this is fixed for each llm call + }, + }, + math_problems=[ + { + "question": problem_inst["question"], + "answer": task.extract_groundtruth(problem_inst["answer"]), + } + ], + llm_gen_fn=lm_call, + ) + + search_tree = CriticSearchTree( + cfg={ + "pb_c_base": config.pb_c_base, + "pb_c_init": config.pb_c_init, + "init_critic_value": config.init_critic_value, + } + ) + rm_call_fn = functools.partial(rm_call, lm_step_tag=lm_call.lm_step_tag) + traj_list = search_tree.critic_mcts( + simulate_env=env, + num_path=config.num_path, + reward_model_fn=rm_call_fn, + select_by_prior=config.select_by_prior + ) + return TreeSearchSolutionOutput( + solutions=[t["text"] for t in traj_list], + completion_tokens=[t["api_completion_tokens"] for t in traj_list], + tree_completion_tokens=[t["tree_completion_tokens"] for t in traj_list], + metadata={ + "review": [[f"Path {i}'s Review"] + t['review_path'] for i, t in enumerate(traj_list)], + "rewrite": [[f"Path {i}'s Rewrite"] + t['rewrite_path'] for i, t in enumerate(traj_list)] + } + ) \ No newline at end of file diff --git a/reason/guided_search/critic_mcts.py b/reason/guided_search/critic_mcts.py new file mode 100644 index 0000000..a13c6bd --- /dev/null +++ b/reason/guided_search/critic_mcts.py @@ -0,0 +1,221 @@ +""" +Critic-guided MCTS Reasoning +""" +import copy +import json +import math + +import numpy as np +import torch +import torch.nn as nn +from typing import List, Dict, Any, Optional, Tuple, Union, Callable, Type +from distributed.utils import print_rank_0, print_with_rank +from envs.base_env import CoTEnv +import pdb +from tqdm import tqdm +import heapq +from copy import deepcopy +from tqdm import tqdm +from collections import defaultdict + +from .tree import Node, LanguageNode, SearchTree +from envs.rstar.rstar_utils import * + +class CriticLanguageNode(LanguageNode): + def __init__(self, from_review: Optional[str] = None, *args, **kwargs): + super().__init__(*args, **kwargs) + self.from_review = from_review + + +class CriticSearchTree(SearchTree): + def __init__(self, cfg) -> None: + super().__init__(cfg) + + self.show_tree_expansion = False + + def critic_mcts( + self, + simulate_env: Type[CoTEnv], + num_path: int, + reward_model_fn: Optional[Callable] = None, + select_by_prior: bool = False, + ) -> List[Dict]: + + api_call_completion_tokens = 0 + _, info = simulate_env.reset(update_legal_action=True) + api_call_completion_tokens += info["api_completion_token"] + if self.root is None: + root = CriticLanguageNode(text_state=simulate_env.get_state()) + self._expand_leaf_node(root, simulate_env, reward_model_fn) + self.root = root + + traj_list = [] + for i_path in range(num_path): + node = self.root + env_copy = simulate_env.copy() + done = False + review_traj = [] + rewrite_traj = [] + + while not done: + if node.visit_count > 0: + action, node = self._select_child(node, env_copy) + else: + if select_by_prior: + action, node = self._select_by_prior(node, env_copy) + else: + action, node = self._select_child(node, env_copy) + + env_copy._next_state_terminated = {} + assert node.last_action == action + env_copy._next_state_terminated[action] = node.terminated + + review_traj.append(node.from_review) + rewrite_traj.append(node.last_action) + + _, _, terminated, truncated, info = env_copy.step( + action, update_legal_action=node.is_leaf() + ) + + done = terminated or truncated + + if not done and node.is_leaf(): + self._expand_leaf_node(node, env_copy, reward_model_fn) + + # record api_tokens, if not expand, info["api_completion_token"] is 0 + api_call_completion_tokens += info["api_completion_token"] + else: + if node.visit_count > 0: + leaf_value = node.value + else: + if self._init_critic_value: + leaf_value = node._initial_value + else: + leaf_value = reward_model_fn(env_copy.get_state()).item() + node.update_recursive(leaf_value, env_copy.mcts_mode) + + traj_data = { + "path_idx": i_path, + "text": env_copy.full_answer, # final answer + "value": leaf_value, + "api_completion_tokens": api_call_completion_tokens, + "tree_completion_tokens": self._completion_tokens, + "review_path": review_traj, + "rewrite_path": rewrite_traj + } + + traj_list.append(traj_data) + + # reset api_call_completion_tokens + api_call_completion_tokens = 0 + + return traj_list + + def _expand_leaf_node( + self, + node: Node, + simulate_env: Type[CoTEnv], + reward_fn: Optional[Callable] = None, + ) -> float: + """ + Overview: + expand the node with the reward_fn. + Arguments: + - node (:obj:`Class Node`): current node when performing mcts search. + - simulate_env (:obj:`Class BaseGameEnv`): the class of simulate env. + - reward_fn (:obj:`Function`): the Callable to compute the state value. + Returns: + - leaf_value (:obj:`Bool`): the leaf node's value. + """ + """ + action_probs_dict, leaf_value = reward_fn(simulate_env) + for action, prior_p in action_probs_dict.items(): + if action in simulate_env.legal_actions: + node.children[action] = Node(parent=node, prior_p=prior_p) + """ + + text_state = simulate_env.get_state() + if not self._init_critic_value: + leaf_value = reward_fn(text_state) + + else: + leaf_value = node._initial_value + assert len(simulate_env.legal_actions) > 0 + prms = reward_fn( + [ + ( + simulate_env.question, + simulate_env.answer + x["action"], + ) + for x in simulate_env.legal_actions + ] + ) + child_values = [] + + # get the last reward + last_r = [i[-1] for i in prms] + + # PRM get last r as single reward + for act, rs in zip(simulate_env.legal_actions, last_r): + # prm-last + child_values.append(rs) + # # prm-min + # child_values.append(min(rs)) + # # prob-prm + # child_values.append(act['prob']) + + assert len(node.children) == 0 + + # for i in simulate_env.legal_actions: + # print(f'\n{i}') + + for i, action_dict in enumerate(simulate_env.legal_actions): + action, prob = action_dict["action"], action_dict["prob"] + + if self._init_critic_value: + child_value = child_values[i] + else: + # XXX(ziyu): consider turn off this branch, i.e. always assume + # `self._init_critic=True`, since with LLM + child_value = 0.0 + + node.children[action] = CriticLanguageNode( + parent=node, + prior_p=prob, + # prm_value=prm_value, + text_state=text_state, + last_action=action, + initial_value=child_value, + num_generated_token=action_dict["num_token"], + from_review=action_dict['from_review'] + ) + # set terminal node here + # print(f"next state teminated = {simulate_env._next_state_terminated}") + # print(f"Action = {action}") + if simulate_env._next_state_terminated[action]: + node.children[action].set_as_terminate_node() + if len(node.children) == 0: + print_rank_0( + "Prune all current children at node {}".format(node.last_action) + ) + + # collect num tokens + if not node.has_collected_token_num: + self._completion_tokens += sum( + c.num_generated_token for c in node.children.values() + ) + node.has_collected_token_num = True + else: + raise RuntimeError("Token number has been collected again.") + + return leaf_value + + def draw_tree(self): + + def display_tree(node): + print("|" + "-" * (node.depth * 4) + str(node)) + for child in node.children: + display_tree(child) + + print(f"\n---------Expanded Tree---------") + display_tree(self.root) \ No newline at end of file diff --git a/reason/guided_search/tree.py b/reason/guided_search/tree.py index 333e220..977cf2d 100644 --- a/reason/guided_search/tree.py +++ b/reason/guided_search/tree.py @@ -395,7 +395,7 @@ def vanila_mcts( traj_data = { "path_idx": i_path, - "text": env_copy.answer, + "text": env_copy.full_answer, "value": leaf_value, "api_completion_tokens": api_call_completion_tokens, "tree_completion_tokens": self._completion_tokens, @@ -549,9 +549,9 @@ def _simulate( if not self.no_terminal_reward: if winner is not None: if winner == 1: - self.answers.add(simulate_env.answer) + self.answers.add(simulate_env.full_answer) else: - self.wrong_answers.add(simulate_env.answer) + self.wrong_answers.add(simulate_env.full_answer) # if simulate_env.mcts_mode == 'self_play_mode': # if winner == -1: @@ -585,7 +585,7 @@ def _simulate( if self.visited_paths is not None: self.visited_paths.append( { - "text": simulate_env.answer, + "text": simulate_env.full_answer, "correct": winner == 1, "value": leaf_value, } @@ -744,7 +744,7 @@ def _expand_leaf_node( print_rank_0( "Prune all current children at node {}".format(node.last_action) ) - + # collect num tokens if not node.has_collected_token_num: self._completion_tokens += sum( diff --git a/scripts/eval/critic_mcts.sh b/scripts/eval/critic_mcts.sh new file mode 100644 index 0000000..8d9fb3b --- /dev/null +++ b/scripts/eval/critic_mcts.sh @@ -0,0 +1,14 @@ +python reason/evaluation/evaluate.py \ + --LM Qwen2.5-Math-1.5B-Instruct \ + --RM math-shepherd-mistral-7b-prm \ + --task_name critic_MATH \ + --max_new_tokens 2048 \ + --num_sequence 16 \ + --tree_max_width 4 \ + --tree_max_depth 4 \ + --save_dir debug_critic_mcts \ + --method critic_mcts \ + --num_worker 32 \ + --controller_addr http://0.0.0.0:28777 \ + --temperature 0.7 + --local