add flexible

lm_eval/tasks/ai2d/ai2d.yaml

@@ -20,6 +20,11 @@ filter_list:
       - function: "regex"
         regex_pattern: "(?:[`\\*_]*(?i:FINAL ANSWER|Final Answer|Answer)[`\\*_]*)[:\\s]*[`\\*_]*([A-D])[`\\*_]*"
       - function: "take_first"
+  - name: "flexible-extract"
+    filter:
+      - function: "custom"
+      - filter_fn: !function utils.flexible_extract
+      - function: "take_first"
 metric_list:
   - metric: exact_match
     aggregation: mean

lm_eval/tasks/ai2d/utils.py

+import string
+
+
+def flexible_extract(resps, docs):
+    def filter_set(inst):
+        filtered = []
+        for resp in inst:
+            while resp[-1] in string.punctuation:
+                resp = resp[:-1]
+            if resp[-1] in ["A", "B", "C", "D"]:
+                resp = resp[-1]
+            filtered.append(resp)
+        return filtered
+
+    filtered_resps = list(map(lambda x: filter_set(x), resps))
+
+    return filtered_resps