A more general solution to model answer extraction instead of output_regex

community_tasks/_template.py

@@ -99,7 +99,7 @@ def __init__(
             suite=["community"],
             generation_size=-1,
             stop_sequence=None,
-            output_regex=None,
+            answer_extractor=None,
             frozen=False,
         )
 

community_tasks/arabic_evals.py

@@ -95,7 +95,7 @@ def __init__(
             suite=["community"],
             generation_size=-1,
             stop_sequence=None,
-            output_regex=None,
+            answer_extractor=None,
             frozen=False,
             trust_dataset=True,
             version=0,
@@ -152,7 +152,7 @@ def __init__(
             suite=["community"],
             generation_size=-1,
             stop_sequence=None,
-            output_regex=None,
+            answer_extractor=None,
             frozen=False,
             trust_dataset=True,
             version=0,
@@ -254,7 +254,7 @@ def __init__(
             suite=["community"],
             generation_size=-1,
             stop_sequence=None,
-            output_regex=None,
+            answer_extractor=None,
             frozen=False,
             version=0,
         )

examples/nanotron/custom_evaluation_tasks.py

@@ -267,7 +267,7 @@ def __init__(
         generation_size=40,
         trust_dataset=True,
         stop_sequence=None,
-        output_regex=None,
+        answer_extractor=None,
         frozen=False,
     ):
         super().__init__(
@@ -282,7 +282,7 @@ def __init__(
             few_shots_select=few_shots_select,
             suite=suite,
             generation_size=generation_size,
-            output_regex=output_regex,
+            answer_extractor=answer_extractor,
             frozen=frozen,
             trust_dataset=trust_dataset,
             stop_sequence=(stop_sequence if stop_sequence is not None else ["\n"]),
@@ -371,7 +371,7 @@ def __init__(
         generation_size=-1,
         trust_dataset=True,
         stop_sequence=None,
-        output_regex=None,
+        answer_extractor=None,
         frozen=False,
     ):
         super().__init__(
@@ -388,7 +388,7 @@ def __init__(
             generation_size=generation_size,
             trust_dataset=trust_dataset,
             stop_sequence=(stop_sequence if stop_sequence is not None else ["\n"]),
-            output_regex=output_regex,
+            answer_extractor=answer_extractor,
             frozen=frozen,
         )
 
@@ -488,7 +488,7 @@ def __init__(
         generation_size=4,
         trust_dataset=True,
         stop_sequence=None,
-        output_regex=None,
+        answer_extractor=None,
         frozen=False,
     ):
         super().__init__(
@@ -505,7 +505,7 @@ def __init__(
             generation_size=generation_size,
             trust_dataset=trust_dataset,
             stop_sequence=(stop_sequence if stop_sequence is not None else ["\n"]),
-            output_regex=output_regex,
+            answer_extractor=answer_extractor,
             frozen=frozen,
         )
 
@@ -624,7 +624,7 @@ def __init__(
         generation_size=-1,
         trust_dataset=True,
         stop_sequence=None,
-        output_regex=None,
+        answer_extractor=None,
         frozen=False,
     ):
         super().__init__(
@@ -641,7 +641,7 @@ def __init__(
             generation_size=generation_size,
             trust_dataset=trust_dataset,
             stop_sequence=(stop_sequence if stop_sequence is not None else ["\n"]),
-            output_regex=output_regex,
+            answer_extractor=answer_extractor,
             frozen=frozen,
         )
 

examples/nanotron/custom_task.py

@@ -82,7 +82,7 @@ def mmlu_anatomy(line):
         generation_size=5,
         metric=[Metrics.loglikelihood_acc_single_token],
         stop_sequence=["\n"],
-        output_regex=None,
+        answer_extractor=None,
         frozen=False,
     ),
     LightevalTaskConfig(
@@ -98,7 +98,7 @@ def mmlu_anatomy(line):
         generation_size=5,
         metric=[Metrics.loglikelihood_acc_single_token],
         stop_sequence=["\n"],
-        output_regex=None,
+        answer_extractor=None,
         frozen=False,
     ),
 ]

src/lighteval/metrics/__init__.py

@@ -20,10 +20,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import re
+from typing import Optional
 
 from lighteval.metrics.metrics import Metric, MetricCategory
-from lighteval.models.model_output import ModelResponse
+from lighteval.models.model_output import AnswerExtractor, ModelResponse
 from lighteval.tasks.requests import Doc
 from lighteval.utils.utils import as_list
 
@@ -89,7 +89,7 @@ def apply_generative_metric(  # noqa: C901
     responses: list[list[ModelResponse]],
     formatted_docs: list[Doc],
     metrics: list[Metric],
-    output_regex: str = None,
+    answer_extractor: Optional[AnswerExtractor] = None,
     max_num_samples: int = 1,
 ):
     outputs = []
@@ -106,8 +106,8 @@ def apply_generative_metric(  # noqa: C901
         preds = []
 
         for pred_raw in preds_raw:
-            if output_regex is not None:
-                pred = next(iter(re.findall(output_regex, pred_raw)), "")
+            if answer_extractor is not None:
+                pred = answer_extractor(pred_raw, formatted_doc.choices)
             else:
                 pred = pred_raw
             preds.append(pred)

src/lighteval/models/model_output.py

@@ -20,8 +20,11 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import random
+import re
+from abc import abstractmethod
 from dataclasses import dataclass, field
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 import torch
 
@@ -81,3 +84,33 @@ class Batch:
     input_lengths: list[int]
     truncated: list[int]
     padded: list[int]
+
+
+class AnswerExtractor:
+    @abstractmethod
+    def __call__(self, result: str, choices: list[str]) -> str:
+        ...
+
+
+class RegexAnswerExtractor(AnswerExtractor):
+    def __init__(
+        self,
+        regex_list: list[re.Pattern | str],
+        fallback: int | Literal["random", "keep", "empty_string"] = "empty_string",
+    ):
+        self.regex_list: list[re.Pattern] = list(map(re.compile, regex_list))
+        self.fallback = fallback
+
+    def __call__(self, result: str, choices: list[str]) -> str:
+        for pattern in self.regex_list:
+            choice = next(iter(re.findall(pattern, result)), "")
+            if choice in choices:
+                return choice
+        if self.fallback == "random":
+            return random.choice(choices)
+        elif self.fallback == "keep":
+            return result
+        elif self.fallback == "empty_string":
+            return ""
+        else:
+            return choices[self.fallback]