[IFEval] Speed up think tag removal (#792)

* [IFEval] Speed up think tag removal

* Refactor and add tests

* Add moar tests

Author: lewtun

src/lighteval/tasks/extended/ifeval/main.py

@@ -20,7 +20,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import re
 
 import numpy as np
 from aenum import extend_enum
@@ -34,6 +33,7 @@
 )
 from lighteval.tasks.lighteval_task import LightevalTaskConfig
 from lighteval.tasks.requests import Doc
+from lighteval.utils.utils import remove_reasoning_tags
 
 
 # Very specific task where there are no precise outputs but instead we test if the format obeys rules
@@ -55,10 +55,15 @@ def ifeval_prompt(line, task_name: str = None):
     "inst_level_loose_acc",
 ]
 
+REASONING_TAG_PAIRS = [
+    ("<think>", "</think>"),
+]
+
 
 def ifeval_metric(predictions: list[str], formatted_doc: Doc, **kwargs) -> dict:
     response = predictions[0]
-    response = re.sub(r"(<think>)?[\s\S]*?<\/think>", "", response)
+    # Remove the reasoning block to avoid false negatives: https://github.com/huggingface/lighteval/issues/790
+    response = remove_reasoning_tags(response, REASONING_TAG_PAIRS)
 
     # Strict instructions
     instruction_list = formatted_doc.specific["instructions_id_list"]

src/lighteval/utils/utils.py

@@ -229,3 +229,28 @@ def download_dataset_worker(
 
 def safe_divide(numerator: np.ndarray, denominator: float, default_value: float = 0.0) -> np.ndarray:
     return np.where(denominator != 0, numerator / denominator, default_value)
+
+
+def remove_reasoning_tags(text: str, tag_pairs: list[tuple[str, str]]) -> str:
+    """Remove all instances of reasoning tag pairs from text.
+
+    See: https://github.com/huggingface/lighteval/issues/790
+
+    Example:
+    >>> text = "<think> Reasoning section </think> Answer section"
+    >>> tag_pairs = [("<think>", "</think>")]
+    >>> remove_reasoning_tags(text, tag_pairs)
+    ' Answer section'
+    """
+    result = text
+
+    for start_tag, end_tag in tag_pairs:
+        while start_tag in result and end_tag in result:
+            start = result.find(start_tag)
+            end = result.find(end_tag, start)
+            if start != -1 and end != -1:
+                result = result[:start] + result[end + len(end_tag) :]
+            else:
+                break
+
+    return result

tests/utils/test_utils.py

@@ -0,0 +1,63 @@
+# MIT License
+
+# Copyright (c) 2024 The HuggingFace Team
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import unittest
+
+from lighteval.utils.utils import remove_reasoning_tags
+
+
+class TestRemoveReasoningTags(unittest.TestCase):
+    def test_remove_reasoning_tags(self):
+        text = "<think> Reasoning section </think> Answer section"
+        tag_pairs = [("<think>", "</think>")]
+        result = remove_reasoning_tags(text, tag_pairs)
+        self.assertEqual(result, " Answer section")
+
+    def test_remove_multiple_tags(self):
+        text = "<think> Reasoning </think> Interlude <think> More reasoning </think> Answer"
+        tag_pairs = [("<think>", "</think>")]
+        result = remove_reasoning_tags(text, tag_pairs)
+        self.assertEqual(result, " Interlude  Answer")
+
+    def test_no_tags(self):
+        text = "No reasoning tags here."
+        tag_pairs = [("<think>", "</think>")]
+        result = remove_reasoning_tags(text, tag_pairs)
+        self.assertEqual(result, "No reasoning tags here.")
+
+    def test_empty_text(self):
+        text = ""
+        tag_pairs = [("<think>", "</think>")]
+        result = remove_reasoning_tags(text, tag_pairs)
+        self.assertEqual(result, "")
+
+    def test_no_opening_tag(self):
+        text = "No opening tag <think> Reasoning section. </think> Answer section"
+        tag_pairs = [("<think>", "</think>")]
+        result = remove_reasoning_tags(text, tag_pairs)
+        self.assertEqual(result, "No opening tag  Answer section")
+
+    def test_no_closing_tag(self):
+        text = "<think> Reasoning section. Answer section"
+        tag_pairs = [("<think>", "</think>")]
+        result = remove_reasoning_tags(text, tag_pairs)
+        self.assertEqual(result, "<think> Reasoning section. Answer section")