Implement MMLU task

datasets.json

@@ -285,7 +285,7 @@
         "parallel": true,
         "translation": "machine",
         "base": "MMLU",
-        "implemented": false,
+        "implemented": true,
         "group": "Multitask Language Understanding"
     },
     {

evals/datasets_/mmlu.py

@@ -1,5 +1,6 @@
-from collections import Counter, defaultdict
 import random
+from collections import Counter, defaultdict
+
 from datasets import get_dataset_config_names, load_dataset
 from joblib.memory import Memory
 from langcodes import Language, standardize_tag
@@ -119,12 +120,30 @@ def print_datasets_analysis():
 
 # print_datasets_analysis()
 
-def load_mmlu(language_bcp_47, i):
-    categories = sorted(list(set(_load_dataset("masakhane/afrimmlu", "eng")["dev"]["subject"])))
-    category = categories[i % len(categories)]
-    random.seed(i)
-    j = random.randint(0, 100)
-    print(j)
+
+def parse_choices(row):
+    if not isinstance(row["choices"], list):
+        row["choices"] = eval(row["choices"])
+    return row
+
+
+def add_choices(row):
+    row["choices"] = [
+        row["option_a"],
+        row["option_b"],
+        row["option_c"],
+        row["option_d"],
+    ]
+    return row
+
+
+def load_mmlu(language_bcp_47, nr):
+    categories = sorted(
+        list(set(_load_dataset("masakhane/afrimmlu", "eng")["dev"]["subject"]))
+    )
+    category = categories[nr % len(categories)]
+    random.seed(nr)
+    i = random.randint(0, 100)
     tags_afrimmlu = {
         standardize_tag(a, macro=True): a
         for a in _get_dataset_config_names("masakhane/afrimmlu")
@@ -140,21 +159,25 @@ def load_mmlu(language_bcp_47, i):
     )
     if language_bcp_47 in tags_afrimmlu:
         ds = _load_dataset("masakhane/afrimmlu", tags_afrimmlu[language_bcp_47])
-        return ds["test"].filter(lambda x: x["subject"] == category)[j]
+        ds = ds.map(parse_choices)
+        examples = ds["dev"].filter(lambda x: x["subject"] == category)
+        task = ds["test"].filter(lambda x: x["subject"] == category)[i]
+        return "masakhane/afrimmlu", examples, task
     elif language_bcp_47 in tags_global_mmlu:
         ds = _load_dataset("CohereForAI/Global-MMLU", tags_global_mmlu[language_bcp_47])
-        def add_choices(split):
-            split["choices"] = list(zip([split["option_a"], split["option_b"], split["option_c"], split["option_d"]]))
-            return split
         ds = ds.map(add_choices)
-        return ds["test"].filter(lambda x: x["subject"] == category)[j]
+        examples = ds["dev"].filter(lambda x: x["subject"] == category)
+        task = ds["test"].filter(lambda x: x["subject"] == category)[i]
+        return "CohereForAI/Global-MMLU", examples, task
     elif language_bcp_47 in tags_okapi:
         ds = _load_dataset(
             "lighteval/okapi_mmlu", language_bcp_47, trust_remote_code=True
         )
-        return ds["test"].filter(lambda x: x["id"] == f"{category}/test/{j}")[0]
+        examples = ds["dev"].filter(lambda x: x["subject"] == category)
+        task = ds["test"].filter(lambda x: x["id"] == f"{category}/test/{i}")[0]
+        return "lighteval/okapi_mmlu", examples, task
     elif language_bcp_47 in tags_mmlux:
         # loading this is more complicated, todo
-        return None
+        return None, None, None
     else:
-        return None
+        return None, None, None

evals/main.py

@@ -12,7 +12,7 @@ from tasks import tasks
 # ===== config =====
 
 n_sentences = 10
-n_languages = 3
+n_languages = 10
 n_models = 3
 
 # ===== run evaluation and aggregate results =====

evals/tasks.py

@@ -5,10 +5,10 @@ import evaluate
 import pandas as pd
 import sentencepiece as spm
 from datasets_.flores import flores_sentences
+from datasets_.mmlu import load_mmlu
 from joblib.memory import Memory
 from languages import languages, script_name
 from models import complete, transcribe
-from datasets import load_dataset, get_dataset_config_names
 
 cache = Memory(location=".cache", verbose=0).cache
 bleu = evaluate.load("bleu")
@@ -187,47 +187,47 @@ async def mlm_and_evaluate(model, language_bcp_47, nr):
     ]
 
 
-
 @cache
 async def mmlu_and_evaluate(model, language_bcp_47, nr):
-    item = data["test"][nr]
+    ds_name, examples, task = load_mmlu(language_bcp_47, nr)
+    if not task:
+        return []
+
     def format_item(item):
-        return f"""{item['question']}
+        return f"""{item["question"]}
         
-        A: {item['option_a']}
-        B: {item['option_b']}
-        C: {item['option_c']}
-        D: {item['option_d']}
+        A: {item["choices"][0]}
+        B: {item["choices"][1]}
+        C: {item["choices"][2]}
+        D: {item["choices"][3]}
         
         A|B|C|D?"""
+
     messages = []
-    for example in data["dev"].select(range(5)):
-        messages += [{"role": "user", "content": format_item(example)}, {"role": "assistant", "content": example["answer"]}]
-    messages += [{"role": "user", "content": format_item(item)}]
+    for example in examples:
+        messages += [
+            {"role": "user", "content": format_item(example)},
+            {"role": "assistant", "content": example["answer"]},
+        ]
+    messages += [{"role": "user", "content": format_item(task)}]
     reply = await complete(
         model=model,
         messages=messages,
         temperature=0,
         max_tokens=1,
     )
-    print(reply.choices[0].message.content.strip())
-    acc = int(reply.choices[0].message.content.strip() == item["answer"])
+    acc = int(reply.choices[0].message.content[:1].strip() == task["answer"])
     return [
         {
             "model": model,
             "bcp_47": language_bcp_47,
             "task": "mmlu",
-            "dataset": ds,
             "metric": "accuracy",
             "score": acc,
             "sentence_nr": nr,
         }
     ]
 
-from asyncio import run
-results = run(mmlu_and_evaluate("gpt-4o-mini", "fr", 0))
-print(results)
-exit()
 
 @cache
 async def transcribe_and_evaluate(model, language_bcp_47, nr):
@@ -260,6 +260,7 @@ async def transcribe_and_evaluate(model, language_bcp_47, nr):
         }
     ]
 
+
 tasks = [
     partial(translate_and_evaluate, mode="from"),
     partial(translate_and_evaluate, mode="to"),

frontend/src/components/DatasetTable.js

@@ -145,7 +145,7 @@ const DatasetTable = ({ data }) => {
         filter
         filterElement={tasksRowFilterTemplate}
         showFilterMatchModes={false}
-        style={{ minWidth: '10rem', maxWidth: '15rem' }}
+        style={{ minWidth: '10rem', maxWidth: '10rem' }}
         body={tasksBodyTemplate}
       />
       <Column