Added v2 fewshot and refactored

Author: matthewreed26

src/inspect_evals/_registry.py

@@ -49,7 +49,7 @@
 from .piqa import piqa
 from .pubmedqa import pubmedqa
 from .race_h import race_h
-from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2
+from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2, sec_qa_v2_5_shot
 from .squad import squad
 from .swe_bench import swe_bench
 from .truthfulqa import truthfulqa

src/inspect_evals/sec_qa/__init__.py

@@ -1,3 +1,3 @@
-from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2
+from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2, sec_qa_v2_5_shot
 
-__all__ = ["sec_qa_v1", "sec_qa_v1_5_shot", "sec_qa_v2"]
+__all__ = ["sec_qa_v1", "sec_qa_v1_5_shot", "sec_qa_v2", "sec_qa_v2_5_shot"]

src/inspect_evals/sec_qa/sec_qa.py

@@ -8,66 +8,40 @@
 from typing import Any
 
 from inspect_ai import Task, task
-from inspect_ai.dataset import Sample, hf_dataset
+from inspect_ai.dataset import Dataset, Sample, hf_dataset
 from inspect_ai.scorer import choice
 from inspect_ai.solver import multiple_choice, system_message
 
 # setup for problem + instructions for providing answer
-SYSTEM_MESSAGE = """
-The following are multiple choice questions about Computer Security.
-""".strip()
-
-SYSTEM_MESSAGE_FEWSHOT = (
-    SYSTEM_MESSAGE
-    + """
-Some examples are provided below.
+SYSTEM_MESSAGE_FEWSHOT = """
+The following are multiple choice questions about Computer Security. Some examples are provided below.
 
 {examples}
 """.strip()
-)
+
+DATASET_PATH = "zefang-liu/secqa"
+DATASET_SUBSET_NAME_V1 = "secqa_v1"
+DATASET_SUBSET_NAME_V2 = "secqa_v2"
+DATASET_TEST_SPLIT = "test"
+DATASET_FEWSHOT_SPLIT = "dev"
 
 
 @task
 def sec_qa_v1() -> Task:
     """Inspect Task implementing the SecQA benchmark v1 0-shot"""
-    # dataset
-    dataset = hf_dataset(
-        path="zefang-liu/secqa",
-        name="secqa_v1",  # aka subset
-        split="dev",
-        sample_fields=record_to_sample,
-        trust=True,
-    )
-
-    # define task
+    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V1, DATASET_TEST_SPLIT)
     return Task(
         dataset=dataset,
-        solver=[system_message(SYSTEM_MESSAGE), multiple_choice()],
+        solver=[multiple_choice()],
         scorer=choice(),
     )
 
 
 @task
 def sec_qa_v1_5_shot() -> Task:
     """Inspect Task implementing the SecQA benchmark v1 5-shot"""
-    fewshot_samples = hf_dataset(
-        path="zefang-liu/secqa",
-        name="secqa_v1",  # aka subset
-        split="dev",
-        sample_fields=record_to_sample,
-        trust=True,
-    )
-
-    # dataset
-    dataset = hf_dataset(
-        path="zefang-liu/secqa",
-        name="secqa_v1",  # aka subset
-        split="val",
-        sample_fields=record_to_sample,
-        trust=True,
-    )
-
-    # define task
+    fewshot_samples = retrieve_hf_dataset(DATASET_SUBSET_NAME_V1, DATASET_FEWSHOT_SPLIT)
+    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V1, DATASET_TEST_SPLIT)
     return Task(
         dataset=dataset,
         solver=[
@@ -87,23 +61,45 @@ def sec_qa_v1_5_shot() -> Task:
 @task
 def sec_qa_v2() -> Task:
     """Inspect Task implementing the SecQA benchmark v2 0-shot"""
-    # dataset
-    dataset = hf_dataset(
-        path="zefang-liu/secqa",
-        name="secqa_v2",  # aka subset
-        split="dev",
-        sample_fields=record_to_sample,
-        trust=True,
+    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V2, DATASET_TEST_SPLIT)
+    return Task(
+        dataset=dataset,
+        solver=[multiple_choice()],
+        scorer=choice(),
     )
 
-    # define task
+
+@task
+def sec_qa_v2_5_shot() -> Task:
+    """Inspect Task implementing the SecQA benchmark v2 5-shot"""
+    fewshot_samples = retrieve_hf_dataset(DATASET_SUBSET_NAME_V2, DATASET_FEWSHOT_SPLIT)
+    dataset = retrieve_hf_dataset(DATASET_SUBSET_NAME_V2, DATASET_TEST_SPLIT)
     return Task(
         dataset=dataset,
-        solver=[system_message(SYSTEM_MESSAGE), multiple_choice()],
+        solver=[
+            system_message(
+                SYSTEM_MESSAGE_FEWSHOT.format(
+                    examples="\n\n".join(
+                        [sample_to_fewshot(sample=sample) for sample in fewshot_samples]
+                    )
+                )
+            ),
+            multiple_choice(),
+        ],
         scorer=choice(),
     )
 
 
+def retrieve_hf_dataset(name: str, split: str) -> Dataset:
+    return hf_dataset(
+        path=DATASET_PATH,
+        name=name,
+        split=split,
+        sample_fields=record_to_sample,
+        trust=True,
+    )
+
+
 def record_to_sample(record: dict[str, Any]) -> Sample:
     return Sample(
         input=record["Question"],
@@ -114,10 +110,12 @@ def record_to_sample(record: dict[str, Any]) -> Sample:
 
 def sample_to_fewshot(sample: Sample) -> str:
     prob_str = f"""QUESTION:\n{sample.input}"""
-    labeled_choices = []
+    labeled_choices = "\n"
     for i, letter_label in enumerate(["A", "B", "C", "D"]):
-        labeled_choice = f"{letter_label}: {sample.choices[i] if sample.choices is not None else None}"
-        labeled_choices.append(labeled_choice)
+        labeled_choices = (
+            labeled_choices
+            + f"{letter_label}: {sample.choices[i] if sample.choices is not None else None}\n"
+        )
     choices_str = f"""CHOICES:{labeled_choices}"""
     ans_str = f"""ANSWER: {sample.target}"""