Skip to content

Commit f72be7f

Browse files
committed
Add unit test for ThinkingReward custom tag
Added test_custom_tag to verify that ThinkingReward correctly uses the custom tag passed in during initialization. The test confirms: - Responses with custom tag get full reward - Responses with default tag get no reward when custom tag is used All 26 ThinkingReward tests passing.
1 parent 93a65b2 commit f72be7f

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

tests/unit_tests/rl/test_thinking_reward.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,19 @@ def test_call_very_long_thinking_block(self):
203203
result = self.reward("prompt", f"<think>{long_content}</think>")
204204
self.assertEqual(result, 1.0)
205205

206+
def test_custom_tag(self):
207+
"""Test that ThinkingReward uses the custom tag passed in."""
208+
# Create reward with custom Japanese tag
209+
custom_tag_reward = ThinkingReward(tag="思考")
210+
211+
# Response with custom tag should get full reward
212+
result = custom_tag_reward("prompt", "<思考>This is my reasoning</思考>")
213+
self.assertEqual(result, 1.0)
214+
215+
# Response with default "think" tag should get no reward
216+
result = custom_tag_reward("prompt", "<think>This is my reasoning</think>")
217+
self.assertEqual(result, 0.0)
218+
206219

207220
if __name__ == "__main__":
208221
unittest.main()

0 commit comments

Comments
 (0)