Merge pull request #3365 from AI-Hypercomputer:rl-utils-tests

PiperOrigin-RevId: 881558474

Author: Google-ML-Automation

tests/unit/rl_utils_test.py

@@ -23,6 +23,11 @@
     reason="tunix (required by evaluate_rl) is not installed GPU",
 )
 
+utils_rl = pytest.importorskip(
+    "maxtext.trainers.post_train.rl.utils_rl",
+    reason="tunix (required by utils_rl) is not installed GPU",
+)
+
 
 def _make_config():
   """Create a minimal config object with the parameters required by score_responses."""
@@ -102,5 +107,22 @@ def test_with_incomplete_reasoning_tags(self):
     self.assertFalse(has_correct_format)
 
 
+class TestExtractHashAnswer(unittest.TestCase):
+  """Tests for utils_rl.extract_hash_answer."""
+
+  @pytest.mark.cpu_only
+  def test_with_hash(self):
+    """Test extraction when #### is present."""
+    self.assertEqual(utils_rl.extract_hash_answer("The answer is #### 42"), "42")
+    self.assertEqual(utils_rl.extract_hash_answer("Some reasoning ####   123.45  "), "123.45")
+    self.assertEqual(utils_rl.extract_hash_answer("####"), "")
+
+  @pytest.mark.cpu_only
+  def test_without_hash(self):
+    """Test extraction when #### is not present."""
+    self.assertIsNone(utils_rl.extract_hash_answer("The answer is 42"))
+    self.assertIsNone(utils_rl.extract_hash_answer(""))
+
+
 if __name__ == "__main__":
   unittest.main()