Merlin-Arthur pipeline (persistent)

src/eval_framework/tasks/benchmarks/squad.py

@@ -238,6 +238,7 @@ def _get_ground_truth(self, item: dict[str, Any]) -> list[str]:
         return item["answers"]["text"]
 
 
+
 class SQuAD2_MA(SQUAD2):
     """SQuAD v2 with the exact system prompt used in MA training"""
 
@@ -301,6 +302,7 @@ def _get_system_prompt_text(self, item: dict[str, Any]) -> str | None:
         return ""
 
 
+
 class SQuAD_OLMES(SQUAD):
     """SQuAD variant matching OLMES implementation."""
 

src/eval_framework/tasks/benchmarks/triviaqa.py

@@ -0,0 +1,78 @@
+import random
+from typing import Any
+
+from eval_framework.metrics.completion.accuracy_completion import AccuracyCompletion
+from eval_framework.metrics.completion.f1 import F1, F1SquadNormalized
+from eval_framework.tasks.base import BaseTask, Language, ResponseType, Sample
+from eval_framework.tasks.dataset_revisions import HF_REVISIONS_LOCKFILE
+
+
+class TRIVIAQA(BaseTask[str]):
+    """Trivia QA dataset: https://huggingface.co/datasets/mandarjoshi/trivia_qa"""
+
+    REVISION_LOCKFILE = HF_REVISIONS_LOCKFILE
+
+    NAME = "TriviaQA"
+    DATASET_PATH = "mandarjoshi/trivia_qa"
+    SAMPLE_SPLIT = "validation"
+    FEWSHOT_SPLIT = "train"
+    RESPONSE_TYPE = ResponseType.COMPLETION
+    METRICS = [AccuracyCompletion, F1]
+    SUBJECTS = ["rc.wikipedia.nocontext"]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Question", "Answer"]
+    LANGUAGE = Language.ENG
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        super().__init__(num_fewshot)
+        self.stop_sequences = ["\n"]
+        self.max_tokens = 400  # the max length of the ground truth is 282 characters while the average is ~16
+        self.rnd_choice_shuffle = random.Random()
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        prompt = f"Question: {item['question'].strip()}\nAnswer:"
+        return prompt
+
+    def _get_fewshot_target_text(self, item: dict[str, Any]) -> str:
+        target = self._get_ground_truth(item)[0]
+        assert target is not None
+        assert isinstance(target, str)
+        return f" {target}"
+
+    def _get_ground_truth(self, item: dict[str, Any]) -> list[str]:
+        return item["answer"]["aliases"]
+
+    def post_process_generated_completion(self, completion_text: str, sample: Sample | None = None) -> str:
+        return completion_text.strip().rstrip(".")
+
+
+class TriviaQA_MA(TRIVIAQA):
+    """TriviaQA with the exact system prompt used in MA training"""
+
+    REVISION_LOCKFILE = HF_REVISIONS_LOCKFILE
+
+    NAME = "TriviaQA_MA"
+    SUBJECTS = ["rc.wikipedia"]
+    UNANSWERABLE_STR = "unanswerable"
+
+    METRICS = [AccuracyCompletion, F1, F1SquadNormalized]
+    PERTURBATION_UNMODIFIABLE_WORDS = ["Question", "Answer", "Context", "unanswerable"]
+
+    def __init__(self, num_fewshot: int = 0) -> None:
+        super().__init__(num_fewshot)
+        self.stop_sequences = []
+        self.max_tokens = 27_000
+
+    def _get_context_text(self, item: dict[str, Any]) -> str:
+        return "\n\n".join(item["entity_pages"]["wiki_context"])
+
+    def _get_system_prompt_text(self, item: dict[str, Any]) -> str | None:
+        return (
+            "You are a helpful assistant and will answer the user's questions carefully, "
+            "logically, accurately and well-reasoned.\n"
+            "Use the given context to answer the question faithfully. Answer only if the "
+            f"answer is present in the given context, otherwise respond with '{self.UNANSWERABLE_STR}' "
+            "if the answer is not present in the context."
+        )
+
+    def _get_instruction_text(self, item: dict[str, Any]) -> str:
+        return f"Context:\n{self._get_context_text(item)}\n\nQuestion:\n{item['question'].strip()}\n"

src/eval_framework/tasks/hf-dataset-revisions.json

@@ -36,6 +36,7 @@
     "google-research-datasets/nq_open": "5dd9790a83002ad084ddeb7c420dc716852c6f28",
     "google/IFEval": "966cd89545d6b6acfd7638bc708b98261ca58e84",
     "jzhang86/de_ifeval": "4f52d847003b3c83cc282e9d296853a24b34b19a",
+    "mandarjoshi/trivia_qa": "0f7faf33a3908546c6fd5b73a660e0f8ff173c2f",
     "math-ai/aime25": "563bb8404243c5f09de6ec262f2db674fe5bce9b",
     "math-ai/aime26": "79037aebdb6580008fb960d17cb21fd3099083e3",
     "nuprl/MultiPL-E": "28441b6024e71d4a1c1c0f6bf171c935cd5a43f2",

src/eval_framework/tasks/task_names.py

@@ -37,6 +37,7 @@ def register_all_tasks(registry: Registry | None = None) -> None:
     register_piqa_tasks(registry=registry)
     register_sciq_tasks(registry=registry)
     register_squad_tasks(registry=registry)
+    register_triviaqa_tasks(registry=registry)
     register_winogrande_tasks(registry=registry)
     register_csqa_tasks(registry=registry)
     register_drop_tasks(registry=registry)
@@ -196,6 +197,12 @@ def register_squad_tasks(registry: Registry) -> None:
     register_lazy_task("eval_framework.tasks.benchmarks.squad.SQuAD2_MA_NO_SYSPROMPT", registry=registry)
 
 
+def register_triviaqa_tasks(registry: Registry) -> None:
+    """Register triviaqa benchmark tasks."""
+    register_lazy_task("eval_framework.tasks.benchmarks.triviaqa.TRIVIAQA", registry=registry)
+    register_lazy_task("eval_framework.tasks.benchmarks.triviaqa.TriviaQA_MA", registry=registry)
+
+
 def register_winogrande_tasks(registry: Registry) -> None:
     """Register winogrande benchmark tasks."""
     register_lazy_task("eval_framework.tasks.benchmarks.winogrande.WINOGRANDECloze", registry=registry)

tests/tests_eval_framework/tasks/benchmarks/task-prompts-hashes.json

@@ -354,6 +354,10 @@
     "SocialIQAMC_OLMES.ConcatFormatter": "484d60e7f583eb6739d6027fec255d06",
     "SocialIQAMC_OLMES.Llama3Formatter": "83a826a85d522f5bc40ecde5599a9dc5",
     "SocialIQAMC_OLMES.NoStripConcatFormatter": "484d60e7f583eb6739d6027fec255d06",
+    "TRIVIAQA.ConcatFormatter": "33595d3c66378ea01b0d1c99a37761e4",
+    "TRIVIAQA.Llama3Formatter": "663385749f5605d86f19b39ce74365eb",
+    "TriviaQA_MA.ConcatFormatter": "ba2f5447024449093a6c07103469d0d8",
+    "TriviaQA_MA.Llama3Formatter": "c53ed23137d14fc5e1d95a711ff83905",
     "WINOGRANDECloze.ConcatFormatter": "a7ff60e76f35d94da1f329a23e0f889b",
     "WINOGRANDECloze.Llama3Formatter": "1c21b8be1a51a76ab39a9be7f7dc0cc4",
     "WINOGRANDECloze.NoStripConcatFormatter": "a7ff60e76f35d94da1f329a23e0f889b",
@@ -366,4 +370,4 @@
     "WINOGRANDE_ELLAMIND_PARTIAL_EVAL_DE.ConcatFormatter": "bb9c64876b5dc3894064a35eabbe29ec",
     "WINOGRANDE_ELLAMIND_PARTIAL_EVAL_DE.Llama3Formatter": "1a679d947f96633962d769eb1139ab38",
     "WINOGRANDE_ELLAMIND_PARTIAL_EVAL_DE.NoStripConcatFormatter": "bb9c64876b5dc3894064a35eabbe29ec"
-}
+}

tests/tests_eval_framework/tasks/benchmarks/test_triviaqa.py

@@ -0,0 +1,45 @@
+import pytest
+
+from eval_framework.tasks.benchmarks.triviaqa import TriviaQA_MA
+from eval_framework.tasks.registry import Registry
+from eval_framework.tasks.task_names import register_triviaqa_tasks
+from template_formatting.formatter import BaseFormatter, ConcatFormatter, Llama3Formatter
+from tests.tests_eval_framework.tasks.benchmarks.utils import run_formatter_hash_test
+
+# Registry for this test suite only holding triviaqa tasks
+_triviaqa_registry = Registry()
+register_triviaqa_tasks(registry=_triviaqa_registry)
+
+
+@pytest.mark.formatter_hash
+@pytest.mark.parametrize("formatter_cls", [Llama3Formatter, ConcatFormatter])
+@pytest.mark.parametrize("task_name", _triviaqa_registry.task_names())
+def test_formatter_hash(task_name: str, formatter_cls: type[BaseFormatter]) -> None:
+    run_formatter_hash_test(task_name, formatter_cls, registry=_triviaqa_registry)
+
+
+@pytest.fixture
+def item():
+    return {
+        "question": "What is the capital of France?",
+        "answer": {"aliases": ["Paris", "Paris, France"]},
+        "entity_pages": {"wiki_context": ["Paris is the capital of France.", "France is in Europe."]},
+    }
+
+
+def test_triviaqa_ma_system_prompt_instructs_reject(item) -> None:
+    system = TriviaQA_MA()._get_system_prompt_text(item)
+    assert system is not None
+    assert f"respond with '{TriviaQA_MA.UNANSWERABLE_STR}'" in system
+
+
+def test_triviaqa_ma_instruction_is_context_question_only(item) -> None:
+    instruction = TriviaQA_MA()._get_instruction_text(item)
+    assert instruction == (
+        "Context:\nParis is the capital of France.\n\nFrance is in Europe.\n\n"
+        "Question:\nWhat is the capital of France?\n"
+    )
+
+
+def test_triviaqa_ma_ground_truth_uses_aliases(item) -> None:
+    assert TriviaQA_MA()._get_ground_truth(item) == ["Paris", "Paris, France"]