fix(rstar): preserve signed and decimal answers during search

Author: liuzhengyang699

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

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)