Merge rewards files

src/forge/data/rewards/math.py → src/forge/data/rewards.py

@@ -54,3 +54,19 @@ def __call__(self, prompt: str, response: str, target: str) -> float:
         if abs(expected_answer - model_answer) < self.tolerance:
             return 1.0  # Correct answer
         return 0.0  # Incorrect answer
+
+
+class ThinkingReward(Reward):
+    """Reward class for evaluating use of <think> tags in reasoning."""
+
+    def __init__(self, reward_value: float = 0.5):
+        self.reward_value = reward_value
+
+    def __call__(
+        self, prompt: str, response: str, target: Optional[str] = None
+    ) -> float:
+        """Check if response contains <think>...</think> tags."""
+        resp = response.lower()
+        if "<think>" in resp and "</think>" in resp:
+            return self.reward_value
+        return 0.0

tests/unit_tests/rl/test_math_reward.py

@@ -7,7 +7,7 @@
 import unittest
 from unittest import mock
 
-from forge.data.rewards.math import MathReward
+from forge.data.rewards import MathReward
 
 
 class TestMathReward(unittest.TestCase):

tests/unit_tests/rl/test_thinking_reward.py

@@ -6,7 +6,7 @@
 
 import unittest
 
-from forge.data.rewards.thinking import ThinkingReward
+from forge.data.rewards import ThinkingReward
 
 
 class TestThinkingReward(unittest.TestCase):