"""

WMDP: Weapons of Mass Destruction Proxy

Nathaniel Li, Alexander Pan, Anjali Gopal, Summer Yue,

Daniel Berrios, Alice Gatti, Justin D Li, Ann-Kathrin

Dombrowski, Shashwat Goel, Long Phan, et al.

https://arxiv.org/pdf/2403.03218

# eval w/ 500 randomly selected samples

inspect eval inspect_evals/wmdp_bio --limit 500

inspect eval inspect_evals/wmdp_chem --limit 500

inspect eval inspect_evals/wmdp_cyber --limit 500

"""

from typing import Any

from inspect_ai import Task, task

from inspect_ai.dataset import Sample

from inspect_ai.scorer import choice

from inspect_ai.solver import multiple_choice

from inspect_evals.metadata import load_eval_metadata

from inspect_evals.utils import create_stable_id, filter_duplicate_ids

from inspect_evals.utils.huggingface import hf_dataset

DATASET_PATH = "cais/wmdp"

WMDP_DATASET_REVISION = "7125571f22f032c56415e7980f48d877dd830ff8"

EVAL_VERSION = load_eval_metadata("wmdp").version

@task

def wmdp_bio() -> Task:

return create_task("wmdp-bio")

@task

def wmdp_chem() -> Task:

return create_task("wmdp-chem")

@task

def wmdp_cyber() -> Task:

return create_task("wmdp-cyber")

def create_task(dataset_name: str) -> Task:

dataset = hf_dataset(

path=DATASET_PATH,

name=dataset_name,

sample_fields=record_to_sample,

split="test",

revision=WMDP_DATASET_REVISION,

)

# Filter out any samples with duplicate IDs (from identical questions)

dataset = filter_duplicate_ids(dataset)

return Task(

dataset=dataset,

solver=multiple_choice(),

scorer=choice(),

version=EVAL_VERSION.comparability_version,

metadata=EVAL_VERSION.to_metadata(),

)

def record_to_sample(record: dict[str, Any]) -> Sample:

choices_map = {0: "A", 1: "B", 2: "C", 3: "D"}

return Sample(

id=create_stable_id(record["question"], record["choices"], prefix="wmdp"),

input=record["question"],

choices=record["choices"],

target=choices_map[record["answer"]],

)