diff --git a/optillm/rstar.py b/optillm/rstar.py index 4b2aa60..a653ede 100644 --- a/optillm/rstar.py +++ b/optillm/rstar.py @@ -317,13 +317,15 @@ def evaluate(self, node: Node) -> float: def extract_answer(self, final_state: str) -> Tuple[str, float]: logger.debug(f"Extracting answer from state: {final_state}") + # Keep the full numeric value when assigning rewards and counting votes. + number_pattern = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?" patterns = [ - r"The answer is (\d+)", - r"The final answer is (\d+)", - r"Therefore, the answer is (\d+)", - r"So, the answer is (\d+)", - r"Thus, the answer is (\d+)", - r"In conclusion, the answer is (\d+)", + rf"The answer is ({number_pattern})", + rf"The final answer is ({number_pattern})", + rf"Therefore, the answer is ({number_pattern})", + rf"So, the answer is ({number_pattern})", + rf"Thus, the answer is ({number_pattern})", + rf"In conclusion, the answer is ({number_pattern})", ] for pattern in patterns: @@ -335,7 +337,7 @@ def extract_answer(self, final_state: str) -> Tuple[str, float]: return answer, confidence # If no pattern is found, try to extract any number - numbers = re.findall(r'\d+', final_state) + numbers = re.findall(number_pattern, final_state) if numbers: answer = numbers[-1] # Take the last number found confidence = 0.5 # Lower confidence as it's not in the expected format diff --git a/tests/test_rstar_numeric_answers.py b/tests/test_rstar_numeric_answers.py new file mode 100644 index 0000000..9ab49e6 --- /dev/null +++ b/tests/test_rstar_numeric_answers.py @@ -0,0 +1,46 @@ +import unittest +from unittest.mock import Mock + +from optillm.rstar import Node, RStar + + +class TestRStarNumericAnswers(unittest.TestCase): + def setUp(self): + self.search = RStar("", Mock(), "mock", max_depth=1, num_rollouts=1) + + def test_explicit_answers_keep_the_complete_number(self): + for answer in ("-3", "+3", "1.25", "-0.25", ".5", "-.5", "1.2e-3", "-2E+4", "42"): + for prefix in ("The answer is", "The final answer is", "Thus, the answer is"): + with self.subTest(answer=answer, prefix=prefix): + self.assertEqual( + self.search.extract_answer(f"{prefix} {answer}."), + (answer, 1.0), + ) + + def test_fallback_keeps_sign_decimal_and_exponent(self): + for answer in ("-3", "1.25", "-.5", "1.2e-3", "42"): + with self.subTest(answer=answer): + self.assertEqual( + self.search.extract_answer(f"After 2 steps, we obtain {answer}."), + (answer, 0.5), + ) + + def test_answer_without_numbers_remains_unknown(self): + self.assertEqual(self.search.extract_answer("More reasoning is needed."), ("", 0.0)) + + def test_signed_answer_receives_explicit_answer_reward(self): + node = Node("The final answer is -1.25", "A1") + self.assertEqual(self.search.evaluate(node), 1.0) + + def test_answer_voting_does_not_merge_different_decimal_values(self): + states = ["The answer is 1.25", "The answer is 1.75", "The answer is 1.75"] + answers = [self.search.extract_answer(state) for state in states] + self.assertEqual(self.search.select_best_answer(answers), "1.75") + + def test_solve_preserves_the_generated_numeric_answer(self): + self.search.generate_response = Mock(return_value="The final answer is -1.25") + + answer, tokens = self.search.solve("Compute the result.") + + self.assertEqual(answer, "-1.25") + self.assertEqual(tokens, 0)