Implemented the v1 5-shot

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_v2
+from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2
 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_v2
+from .sec_qa import sec_qa_v1, sec_qa_v1_5_shot, sec_qa_v2
 
-__all__ = ["sec_qa_v1", "sec_qa_v2"]
+__all__ = ["sec_qa_v1", "sec_qa_v1_5_shot", "sec_qa_v2"]

src/inspect_evals/sec_qa/sec_qa.py

@@ -17,10 +17,19 @@
 The following are multiple choice questions about Computer Security.
 """.strip()
 
+SYSTEM_MESSAGE_FEWSHOT = (
+    SYSTEM_MESSAGE
+    + """
+Some examples are provided below.
+
+{examples}
+""".strip()
+)
+
 
 @task
 def sec_qa_v1() -> Task:
-    """Inspect Task implementing the SecQA benchmark v1"""
+    """Inspect Task implementing the SecQA benchmark v1 0-shot"""
     # dataset
     dataset = hf_dataset(
         path="zefang-liu/secqa",
@@ -38,9 +47,46 @@ def sec_qa_v1() -> Task:
     )
 
 
+@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
+    return Task(
+        dataset=dataset,
+        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(),
+    )
+
+
 @task
 def sec_qa_v2() -> Task:
-    """Inspect Task implementing the SecQA benchmark v2"""
+    """Inspect Task implementing the SecQA benchmark v2 0-shot"""
     # dataset
     dataset = hf_dataset(
         path="zefang-liu/secqa",
@@ -64,3 +110,15 @@ def record_to_sample(record: dict[str, Any]) -> Sample:
         choices=[record["A"], record["B"], record["C"], record["D"]],
         target=record["Answer"],
     )
+
+
+def sample_to_fewshot(sample: Sample) -> str:
+    prob_str = f"""QUESTION:\n{sample.input}"""
+    labeled_choices = []
+    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)
+    choices_str = f"""CHOICES:{labeled_choices}"""
+    ans_str = f"""ANSWER: {sample.target}"""
+
+    return f"""{prob_str}\n\n{choices_str}\n\n{ans_str}"""