add task variants

lm_eval/tasks/mgsm/mgsm_th_direct.yaml

+# Generated by utils.py
+dataset_name: th
+doc_to_target: '{% if answer is not none %}{{answer[6+1]}}{% else %}{{answer_number|string}}{%
+  endif %}'
+doc_to_text: '{% if answer is not none %}{{question+"\nAnswer"}}{% else %}{{"โจทย์:
+  "+question+"\nAnswer"}}{% endif %}'
+include: direct_yaml
+task: mgsm_th_direct

lm_eval/tasks/mgsm/mgsm_th.yaml → lm_eval/tasks/mgsm/mgsm_th_native-cot.yaml

@@ -4,5 +4,11 @@ doc_to_target: '{% if answer is not none %}{{answer[17+1]}}{% else %}{{answer_nu
   endif %}'
 doc_to_text: '{% if answer is not none %}{{question+"\nคำตอบทีละขั้นตอน:"}}{% else
   %}{{"โจทย์: "+question+"\nคำตอบทีละขั้นตอน:"}}{% endif %}'
-include: common_template_yaml
-task: mgsm_th
+filter:
+- function: regex
+  regex_pattern: The answer is (\-?[0-9\.\,]+)
+- function: take_first
+filter_list:
+- name: get-answer
+include: cot_yaml
+task: mgsm_th_direct

lm_eval/tasks/mgsm/mgsm_zh_direct.yaml

+# Generated by utils.py
+dataset_name: zh
+doc_to_target: '{% if answer is not none %}{{answer[6+1]}}{% else %}{{answer_number|string}}{%
+  endif %}'
+doc_to_text: '{% if answer is not none %}{{question+"\nAnswer"}}{% else %}{{"问题: "+question+"\nAnswer"}}{%
+  endif %}'
+include: direct_yaml
+task: mgsm_zh_direct

lm_eval/tasks/mgsm/mgsm_zh.yaml → lm_eval/tasks/mgsm/mgsm_zh_native-cot.yaml

@@ -4,5 +4,11 @@ doc_to_target: '{% if answer is not none %}{{answer[5+1]}}{% else %}{{answer_num
   endif %}'
 doc_to_text: '{% if answer is not none %}{{question+"\n逐步解答:"}}{% else %}{{"问题: "+question+"\n逐步解答:"}}{%
   endif %}'
-include: common_template_yaml
-task: mgsm_zh
+filter:
+- function: regex
+  regex_pattern: The answer is (\-?[0-9\.\,]+)
+- function: take_first
+filter_list:
+- name: get-answer
+include: cot_yaml
+task: mgsm_zh_direct

lm_eval/tasks/mgsm/utils.py

@@ -6,51 +6,94 @@ LANGUAGES = {
     "bn": {  # Bengali
         "QUESTION": "\u09aa\u09cd\u09b0\u09b6\u09cd\u09a8:",
         "ANSWER": "\u09a7\u09be\u09aa\u09c7 \u09a7\u09be\u09aa\u09c7 \u0989\u09a4\u09cd\u09a4\u09b0:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "de": {  # German
         "QUESTION": "Frage:",
         "ANSWER": "Schritt-f\u00fcr-Schritt-Antwort:",
+        "DIRECT": "Antwort:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "en": {  # English
         "QUESTION": "Question:",
         "ANSWER": "Step-by-Step Answer:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "es": {  # Spanish
         "QUESTION": "Pregunta:",
         "ANSWER": "Respuesta paso a paso:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "fr": {  # French
         "QUESTION": "Question :",
         "ANSWER": "R\u00e9ponse \u00e9tape par \u00e9tape :",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "ru": {  # Russian
         "QUESTION": "\u0417\u0430\u0434\u0430\u0447\u0430:",
         "ANSWER": "\u041f\u043e\u0448\u0430\u0433\u043e\u0432\u043e\u0435\u0440\u0435\u0448\u0435\u043d\u0438\u0435:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "sw": {  # Swahili
         "QUESTION": "Swali:",
         "ANSWER": "Jibu la Hatua kwa Hatua:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "te": {  # Telugu
         "QUESTION": "\u0c2a\u0c4d\u0c30\u0c36\u0c4d\u0c28:",
         "ANSWER": "\u0c26\u0c36\u0c32\u0c35\u0c3e\u0c30\u0c40\u0c17\u0c3e \u0c38\u0c2e\u0c3e\u0c27\u0c3e\u0c28\u0c02:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "th": {  # Thai
         "QUESTION": "\u0e42\u0e08\u0e17\u0e22\u0e4c:",
         "ANSWER": "\u0e04\u0e33\u0e15\u0e2d\u0e1a\u0e17\u0e35\u0e25\u0e30\u0e02\u0e31\u0e49\u0e19\u0e15\u0e2d\u0e19:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "ja": {  # Japanese
         "QUESTION": "\u554f\u984c:",
         "ANSWER": "\u30b9\u30c6\u30c3\u30d7\u3054\u3068\u306e\u7b54\u3048:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
     "zh": {  # Chinese
         "QUESTION": "\u95ee\u9898:",
         "ANSWER": "\u9010\u6b65\u89e3\u7b54:",
+        "DIRECT": "Answer:",
+        "REGEX": "The answer is (\\-?[0-9\\.\\,]+)",
     },
 }
 
+def add_regex_pattern(regex_pattern):
 
-def gen_lang_yamls(output_dir: str, overwrite: bool) -> None:
+    if regex_pattern is None:
+        return {}
+    return {
+                "filter_list": [
+                    {
+                        "name": "get-answer",
+                        },
+                    ],
+                "filter": [
+                    {
+                        "function": "regex",
+                        "regex_pattern": regex_pattern,
+                        },
+                    {
+                        "function": "take_first",
+                        },
+                    ],
+                }
+
+
+def gen_lang_yamls(output_dir: str, overwrite: bool, mode: str) -> None:
     """
     Generate a yaml file for each language.
 
@@ -59,19 +102,36 @@ def gen_lang_yamls(output_dir: str, overwrite: bool) -> None:
     """
     err = []
     for lang in LANGUAGES.keys():
-        file_name = f"mgsm_{lang}.yaml"
         try:
             QUESTION = LANGUAGES[lang]["QUESTION"]
+
+            yaml_template = "cot_yaml"
+            if mode == "direct":
+                ANSWER = LANGUAGES[lang]["DIRECT"]
+                REGEX = None
+                task_name = f"mgsm_{lang}_direct"
+                yaml_template = "direct_yaml"
+            elif mode == "native-cot":
                 ANSWER = LANGUAGES[lang]["ANSWER"]
+                REGEX = LANGUAGES[lang]["REGEX"]
+                task_name = f"mgsm_{lang}_native-cot"
+            elif model == "en-cot":
+                ANSWER = LANGUAGES["en"]["ANSWER"]
+                REGEX = LANGUAGES["en"]["REGEX"]
+                task_name = f"mgsm_{lang}_en-cot"
+
+            file_name = f"{file_name}.yaml"
+            filter_list = add_regex_pattern(REGEX)
+
             with open(
                 f"{output_dir}/{file_name}", "w" if overwrite else "x", encoding="utf8"
             ) as f:
                 f.write("# Generated by utils.py\n")
                 yaml.dump(
                     {
-                        "include": "common_template_yaml",
+                        "include": yaml_template,
                         "dataset_name": lang,
-                        "task": f"mgsm_{lang}",
+                        "task": f"mgsm_{lang}_direct",
                         "doc_to_text": f"""{{% if answer is not none %}}""" \
                                        f"""{{{{question+"\\n{ANSWER}"}}}}""" \
                                        f"""{{% else %}}""" \
@@ -82,6 +142,7 @@ def gen_lang_yamls(output_dir: str, overwrite: bool) -> None:
                                          f"""{{% else %}}""" \
                                          f"""{{{{answer_number|string}}}}""" \
                                          f"""{{% endif %}}""",
+                        **filter_list
                     },
                     f,
                     allow_unicode=True,
@@ -108,9 +169,12 @@ def main() -> None:
     parser.add_argument(
         "--output-dir", default=".", help="Directory to write yaml files to"
     )
+    parser.add_argument(
+        "--mode", default="native-cot", choices=["direct", "native-cot", "en-cot"], help="Mode of chain-of-thought"
+    )
     args = parser.parse_args()
 
-    gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite)
+    gen_lang_yamls(output_dir=args.output_dir, overwrite=args.overwrite, mode=args.mode)
 
 
 if __name__ == "__main__":