make extraction better

lm_eval/tasks/ai2d/ai2d.yaml

@@ -18,7 +18,7 @@ filter_list:
   - name: "strict-match"
     filter:
       - function: "regex"
-        regex_pattern: "(?:[`\\*_]*(?i:FINAL ANSWER|Final Answer|Answer)[`\\*_]*)[:\\s]*[`\\*_]*([A-D])[`\\*_]*"
+        regex_pattern: "(?:[`\\*_]*(?i:FINAL ANSWER|Final Answer|Answer|answer is)[`\\*_]*)[:\\s]*[`\\*_]*([A-D])[`\\*_]*"
       - function: "take_first"
   - name: "flexible-extract"
     filter:

lm_eval/tasks/ai2d/utils.py

+import re
 import string
 
 
+REGEX = re.compile(
+    "[`*_]*(?i:FINAL ANSWER|Final Answer|Answer|answer is)[`*_]*[:\s]*[`*_]*([A-D])[`*_]*"
+)
+
+
 def flexible_extract(resps, docs):
     def filter_set(inst):
         filtered = []
         for resp in inst:
+            # first, we try to match the regex pattern
+            if match := REGEX.findall(resp):
+                match = match[-1]
+                if match:
+                    return match
+            # if we can't match the regex pattern, we try to match the last character
             while resp[-1] in string.punctuation:
                 resp = resp[:-1]
             if resp[-1] in ["A", "B", "C", "D"]:
                 resp = resp[-1]
+            else:
+                # match on A-D after a colon (last match), for example option: A.
+                pattern = r":\s*([A-D])"
+                matches = re.findall(pattern, resp)
+                if matches:
+                    resp = matches[-1]
             filtered.append(resp)
         return filtered