from inspect_ai import Task, task
from inspect_ai.dataset import Sample
from inspect_ai.scorer import includes
from inspect_ai.solver import generate

EXAMPLES = "The capital of Italy is Rome.\nThe capital of Japan is Tokyo.\n"


@task
def capitals():
    return Task(
        dataset=[
            Sample(input=EXAMPLES + "The capital of France is", target="Paris"),
            Sample(input=EXAMPLES + "The capital of Argentina is", target="Buenos Aires"),
            Sample(input=EXAMPLES + "The capital of Canada is", target="Ottawa"),
        ],
        solver=generate(),
        scorer=includes(),
    )