lm-evaluation-harness/lm_eval/tasks/metabench/process_docs.py

import hashlib
import re

import datasets


def hash_string(string: str) -> str:
    return hashlib.sha256(string.encode("utf-8")).hexdigest()


def process_arc(dataset: datasets.Dataset) -> datasets.Dataset:
    def _subprocess(doc):
        long_prompt = ""
        for shot in range(1, 26):
            question = doc[f"arc_question_shot_{shot}"]
            doc.pop(f"arc_question_shot_{shot}")
            answer_lab = doc[f"arc_answerKey_shot_{shot}"]
            doc.pop(f"arc_answerKey_shot_{shot}")
            answer_idx = doc[f"arc_choices_shot_{shot}"]["label"].index(answer_lab)
            answer = doc[f"arc_choices_shot_{shot}"]["text"][answer_idx]
            doc.pop(f"arc_choices_shot_{shot}")
            doc.pop(f"arc_idx_shot_{shot}")

            long_prompt = f"{long_prompt}Question: {question}\nAnswer: {answer}\n\n"  # no choices are provided in the few-shot setting (per lines 602-610 of lm_eval.api.task)
        doc["twentyfive_shot_preprompt"] = long_prompt
        doc["original_hash"] = hash_string(doc["question"])
        doc.pop("alltwentyfiveshot_longprompt")
        return doc

    return dataset.map(_subprocess)


def process_gsm8k(dataset: datasets.Dataset) -> datasets.Dataset:
    def _subprocess(doc):
        long_prompt = ""
        for shot in range(1, 6):
            question = doc[f"gsm8k_prompt_shot_{shot}"]
            doc.pop(f"gsm8k_prompt_shot_{shot}")
            answer = doc[f"gsm8k_answer_shot_{shot}"]
            doc.pop(f"gsm8k_answer_shot_{shot}")
            doc.pop(f"gsm8k_idx_shot_{shot}")

            long_prompt = f"{long_prompt}Question: {question}\nAnswer: {answer}\n\n"  # no choices are provided in the few-shot setting (per lines 602-610 of lm_eval.api.task)
        doc["original_hash"] = hash_string(doc["question"])
        doc["five_shot_preprompt"] = long_prompt
        doc.pop("allfiveshot_longprompt")
        return doc

    return dataset.map(_subprocess)


def process_hellaswag(dataset: datasets.Dataset) -> datasets.Dataset:
    def process_txt(text):  # mirrored from hellaswag task
        text = text.strip()
        # NOTE: Brackets are artifacts of the WikiHow dataset portion of HellaSwag.
        text = text.replace(" [title]", ". ")
        text = re.sub("\\[.*?\\]", "", text)
        text = text.replace("  ", " ")
        return text

    def _preprocess(doc):
        ctx = doc["ctx_a"] + " " + doc["ctx_b"].capitalize()
        doc.pop("ctx_a")
        doc.pop("ctx_b")
        doc.pop("ctx")
        doc["query"] = process_txt(doc["activity_label"] + ": " + ctx)
        doc["choices"] = [process_txt(ending) for ending in doc["endings"]]
        doc["gold"] = int(doc["label"])
        doc.pop("activity_label")
        doc.pop("endings")

        long_prompt = ""
        for shot in range(1, 11):
            ctx = (
                doc[f"hellaswag_ctx_a_shot_{shot}"]
                + " "
                + doc[f"hellaswag_ctx_b_shot_{shot}"].capitalize()
            )
            doc.pop(f"hellaswag_ctx_a_shot_{shot}")
            doc.pop(f"hellaswag_ctx_b_shot_{shot}")
            doc.pop(f"hellaswag_ctx_shot_{shot}")
            question = process_txt(
                doc[f"hellaswag_activity_labels_shot_{shot}"] + ": " + ctx
            )
            ending = process_txt(
                doc[f"hellaswag_endings_shot_{shot}"][
                    int(doc[f"hellaswag_label_shot_{shot}"])
                ]
            )
            doc.pop(f"hellaswag_activity_labels_shot_{shot}")
            doc.pop(f"hellaswag_endings_shot_{shot}")
            doc.pop(f"hellaswag_label_shot_{shot}")

            long_prompt = f"{long_prompt}{question} {ending}\n\n"

            doc.pop(f"hellaswag_ind_shot_{shot}")
            doc.pop(f"hellaswag_source_id_shot_{shot}")
            doc.pop(f"hellaswag_split_shot_{shot}")
            doc.pop(f"hellaswag_split_type_shot_{shot}")

        doc["original_hash"] = hash_string(doc["query"])
        doc["ten_shot_preprompt"] = long_prompt
        doc.pop("alltenshot_longprompt")
        return doc

    return dataset.map(_preprocess)


def process_mmlu(dataset: datasets.Dataset) -> datasets.Dataset:
    def _subprocess(doc):
        choices = ["A", "B", "C", "D"]
        long_prompt = f"The following are multiple choice questions (with answers) about {' '.join(doc['subject'].split('_'))}.\n\n"
        for shot in range(1, 6):
            question = doc[f"mmlu_question_shot_{shot}"].strip()
            doc.pop(f"mmlu_question_shot_{shot}")
            answer = choices[int(doc[f"mmlu_answers_shot_{shot}"])]
            choice_A = doc[f"mmlu_choices_shot_{shot}"][0]
            choice_B = doc[f"mmlu_choices_shot_{shot}"][1]
            choice_C = doc[f"mmlu_choices_shot_{shot}"][2]
            choice_D = doc[f"mmlu_choices_shot_{shot}"][3]

            doc.pop(f"mmlu_choices_shot_{shot}")
            doc.pop(f"mmlu_answers_shot_{shot}")
            doc.pop(f"mmlu_ind_shot_{shot}")

            long_prompt = f"{long_prompt}{question}\nA. {choice_A}\nB. {choice_B}\nC. {choice_C}\nD. {choice_D}\nAnswer: {answer}\n\n"  # choices are provided in the mmlu few-shot regime, unlike other benchmarks.

        doc["original_hash"] = hash_string(doc["question"])
        doc["five_shot_preprompt"] = long_prompt
        doc.pop("allfiveshot_longprompt")
        return doc

    return dataset.map(_subprocess)


def process_truthfulqa(dataset: datasets.Dataset) -> datasets.Dataset:
    def _subprocess(doc):
        doc["original_hash"] = hash_string(doc["question"])
        return doc

    return dataset.map(_subprocess)


def process_winogrande(dataset: datasets.Dataset) -> datasets.Dataset:
    def _subprocess(doc):
        long_prompt = ""
        for shot in range(1, 6):
            if doc[f"winogrande_answer_shot_{shot}"] == "1":
                answer = doc[f"winogrande_option1_shot_{shot}"]
            elif doc[f"winogrande_answer_shot_{shot}"] == "2":
                answer = doc[f"winogrande_option2_shot_{shot}"]
            else:
                raise ValueError("Answer not recognised.")

            question = doc[f"winogrande_prompt_shot_{shot}"].replace("_", answer)

            doc.pop(f"winogrande_prompt_shot_{shot}")
            doc.pop(f"winogrande_answer_shot_{shot}")
            doc.pop(f"winogrande_idx_shot_{shot}")
            doc.pop(f"winogrande_option1_shot_{shot}")
            doc.pop(f"winogrande_option2_shot_{shot}")

            long_prompt = f"{long_prompt}{question}\n\n"
        sentence = doc["sentence"]
        doc["original_hash"] = hash_string(doc["sentence"])
        doc["sentence"] = f"{long_prompt}{sentence}"
        doc.pop("allfiveshot_longprompt")
        return doc

    return dataset.map(_subprocess)


def winogrande_doc_to_text(doc):  # Mirrored from the winogrande task
    answer_to_num = {"1": 0, "2": 1}
    return answer_to_num[doc["answer"]]


def winogrande_doc_to_target(doc):  # Mirrored from the winogrande task
    idx = doc["sentence"].index("_") + 1
    return doc["sentence"][idx:].strip()


def winogrande_doc_to_choice(doc):  # Mirrored from the winogrande task
    idx = doc["sentence"].index("_")
    options = [doc["option1"], doc["option2"]]
    return [doc["sentence"][:idx] + opt for opt in options]