Upload from GitHub Actions: TruthfulQA translation WIP

evals/datasets_/arc.py

@@ -54,6 +54,7 @@ def load_uhura_arc_easy(language_bcp_47, nr):
         ds = ds.rename_column("answerKey", "answer")
         train_ids = common_ids_train[nr:nr+3]
         examples = ds["train"].filter(lambda x: x["id"] in train_ids)
+        # raise Exception(language_bcp_47)
         task = ds["test"].filter(lambda x: x["id"] == common_ids_test[nr])[0]
         return "fair-forward/arc-easy-autotranslated", examples, task
     else:

evals/datasets_/truthfulqa.py

@@ -3,6 +3,13 @@ from collections import Counter, defaultdict
 
 from langcodes import Language, standardize_tag
 from rich import print
+from tqdm import tqdm
+import asyncio
+from tqdm.asyncio import tqdm_asyncio
+import os
+
+from datasets import Dataset, load_dataset
+from models import translate_google, google_supported_languages
 
 from datasets_.util import _get_dataset_config_names, _load_dataset
 
@@ -28,3 +35,68 @@ def load_truthfulqa(language_bcp_47, nr):
         return "masakhane/uhura-truthfulqa", examples, task
     else:
         return None, None, None
+
+
+
+def translate_truthfulqa(languages):
+    human_translated = [*tags_uhura_truthfulqa.keys()]
+    untranslated = [
+        lang
+        for lang in languages["bcp_47"].values[:100]
+        if lang not in human_translated and lang in google_supported_languages
+    ]
+    n_samples = 10
+
+    slug = "fair-forward/truthfulqa-autotranslated"
+    for lang in tqdm(untranslated):
+        # check if already exists on hub
+        try:
+            ds_lang = load_dataset(slug, lang)
+        except (ValueError, Exception):
+            print(f"Translating {lang}...")
+            for split in ["train", "test"]:
+                ds = _load_dataset(slug_uhura_truthfulqa, tags_uhura_truthfulqa["en"], split=split)
+                samples = []
+                if split == "train":
+                    samples.extend(ds)
+                else:
+                    for i in range(n_samples):
+                        task = ds[i]
+                        samples.append(task)
+                questions_tr = [
+                    translate_google(s["question"], "en", lang) for s in samples
+                ]
+                questions_tr = asyncio.run(tqdm_asyncio.gather(*questions_tr))
+                choices_texts_concatenated = []
+                for s in samples:
+                    for choice in eval(s["choices"]):
+                        choices_texts_concatenated.append(choice)
+                choices_tr = [
+                    translate_google(c, "en", lang) for c in choices_texts_concatenated
+                ]
+                choices_tr = asyncio.run(tqdm_asyncio.gather(*choices_tr))
+                # group into chunks of 4
+                choices_tr = [
+                    choices_tr[i : i + 4] for i in range(0, len(choices_tr), 4)
+                ]
+
+                ds_lang = Dataset.from_dict(
+                    {
+                        "subject": [s["subject"] for s in samples],
+                        "question": questions_tr,
+                        "choices": choices_tr,
+                        "answer": [s["answer"] for s in samples],
+                    }
+                )
+                ds_lang.push_to_hub(
+                    slug,
+                    split=split,
+                    config_name=lang,
+                    token=os.getenv("HUGGINGFACE_ACCESS_TOKEN"),
+                )
+                ds_lang.to_json(
+                    f"data/translations/mmlu/{lang}_{split}.json",
+                    lines=False,
+                    force_ascii=False,
+                    indent=2,
+                )

evals/plots.py

@@ -45,7 +45,7 @@ pivot_df = pivot_df[[task for task in ordered_tasks if task in pivot_df.columns]
 correlation_matrix = pivot_df.corr()
 
 # Create the correlation plot
-plt.figure(figsize=(12, 10))
+plt.figure(figsize=(8, 6))
 # Create mask for upper triangle including diagonal to show only lower triangle  
 mask = np.triu(np.ones_like(correlation_matrix, dtype=bool))
 
@@ -53,7 +53,7 @@ mask = np.triu(np.ones_like(correlation_matrix, dtype=bool))
 sns.heatmap(
     correlation_matrix, 
     annot=True, 
-    cmap='coolwarm', 
+    cmap='Blues', 
     center=0,
     square=True,
     mask=mask,
@@ -61,7 +61,6 @@ sns.heatmap(
     fmt='.3f'
 )
 
-plt.title('Task Performance Correlation Matrix', fontsize=16, fontweight='bold')
 plt.xlabel('Tasks', fontsize=12)
 plt.ylabel('Tasks', fontsize=12)
 plt.xticks(rotation=45, ha='right')