from typing import Any

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

TEMPLATE = r"""
Answer Yes if the following two sentences are paraphrases. If they are not, answer No.
Do not give any other answer other than Yes or No in your response.
{prompt}
"""


@task
def paws() -> Task:
    """Inspect Task implementation of the PAWS benchmark"""
    dataset = hf_dataset(
        path="paws",
        name="labeled_final",
        sample_fields=record_to_sample,
        split="test",
        auto_id=True,
        shuffle=True,
    )

    return Task(
        dataset=dataset,
        solver=[prompt_template(template=TEMPLATE), generate()],
        scorer=includes(),
    )


def record_to_sample(record: dict[str, Any]) -> Sample:
    sentence1 = f"Sentence1: {record['sentence1']}"
    sentence2 = f"Sentence2: {record['sentence2']}"
    output = "Yes" if record["label"] == 1 else "No"

    return Sample(
        input=f"{sentence1}\n{sentence2}",
        target=output,
        id=record["id"],
    )
