LM-Polygraph/wmt19

1from functools import partial2from tqdm import tqdm3import datasets4 5 6TOP_K = 47 8 9def prepare_base(10    dataset, input_column, output_column, prompt=None11) -> tuple[list[str], list[str], dict[str, list[str]]]:12    x, y = dataset[input_column], dataset[output_column]13    if prompt:14        for i in range(len(x)):15            x[i] = prompt.format(text=x[i])16    return x, y17 18 19def prepare_babi_qa(dataset, input_column, output_column, prompt):20    x, y = [], []21    for inst in dataset:22        inst = inst["story"]23        context = ""24        for text, answer in zip(inst[input_column], inst[output_column]):25            if answer == "":26                context += text + " "27            else:28                x.append(prompt.format(context=context.strip(), question=text))29                y.append(answer)30    return x, y31 32 33def prepare_coqa(34    dataset, input_column, output_column, description, prompt, few_shot_prompt, instruct35):36    def doc_to_text(doc, prompt, i=0):37        # Given a passage p, the conversation history {q1, a1, . . . qi−1, ai−1}38        # and a question qi, the task is to predict the answer ai39        doc_text = ""40        for q, a in zip(doc["questions"][:i], doc["answers"]["input_text"][:i]):41            doc_text += "\n\n" + prompt.format(question=q, answer=a, topk=TOP_K)42        return doc_text43 44    x, y = [], []45    for inst in dataset:46        formatted_description = description.format(story=inst["story"], topk=TOP_K)47        for j, (question, answer) in enumerate(48            zip(inst[input_column], inst[output_column]["input_text"])49        ):50            if instruct:51                assert (52                    few_shot_prompt is not None53                ), "separate few_shot_prompt must be provided for instruction mode."54                few_shot_section = doc_to_text(inst, few_shot_prompt, j)55                if few_shot_section != "":56                    few_shot_section = (57                        "\n\nHere are a few examples of questions and answers:"58                        + few_shot_section59                        + "\n\nNow answer the following question in the same format.\n\n"60                    )61                else:62                    few_shot_section = "\n\n"63            else:64                few_shot_section = doc_to_text(inst, prompt, j) + "\n\n"65            formatted_prompt = (66                formatted_description67                + few_shot_section68                + prompt.format(69                    question=question,70                    answer="",71                )72            )73            x.append(formatted_prompt)74            y.append(answer)75    return x, y76 77 78def prepare_mmlu(79    dataset,80    output_column,81    prompt,82    description,83    mmlu_max_subject_size,84    n_shot,85    few_shot_dataset_func,86    few_shot_prompt,87    instruct,88):89    import numpy as np90    np.random.seed(1)91 92    few_shot_dataset = few_shot_dataset_func()93 94    answers = ["A", "B", "C", "D"]95    subjects = np.array(dataset["subject"])96    few_shot_subjects = np.array(few_shot_dataset["subject"])97    x, y = [], []98    for subject in np.unique(subjects):99        formatted_description = description.format(subject=subject.replace("_", " "), topk=TOP_K)100        if n_shot > 0:101            few_shot_subject = few_shot_dataset.select(102                np.argwhere(few_shot_subjects == subject).flatten()103            )104            few_shot_ids = np.random.choice(105                len(few_shot_subject), n_shot, replace=False106            )107            few_shot_data = few_shot_subject.select(few_shot_ids)108            if instruct:109                assert (110                    few_shot_prompt is not None111                ), "separate few_shot_prompt must be provided for instruction mode."112                formatted_few_shot_prompt = (113                    "Here are a few examples of questions and answers:\n\n"114                )115                for inst in few_shot_data:116                    formatted_few_shot_prompt += (117                        few_shot_prompt.format(118                            choices=inst["choices"],119                            question=inst["question"].strip(),120                            answer=answers[inst["answer"]],121                            topk=TOP_K,122                        )123                        + "\n\n"124                    )125                formatted_few_shot_prompt += (126                    "Now answer the following question in the same format:\n\n"127                )128            else:129                formatted_few_shot_prompt = ""130                for inst in few_shot_data:131                    formatted_few_shot_prompt += (132                        prompt.format(133                            choices=inst["choices"],134                            question=inst["question"].strip(),135                            answer=answers[inst["answer"]],136                        )137                        + "\n"138                    )139 140        subject_data = dataset.select(np.argwhere(subjects == subject).flatten())141 142        if len(subject_data) > mmlu_max_subject_size:143            subject_data = subject_data.select(range(mmlu_max_subject_size))144 145        for inst in subject_data:146            formatted_prompt = prompt.format(147                choices=inst["choices"],148                question=inst["question"].strip(),149                answer="",150            )151            x.append(152                formatted_description153                + "\n\n"154                + formatted_few_shot_prompt155                + formatted_prompt156            )157            y.append(answers[inst[output_column]])158    return x, y159 160 161def prepare_person(dataset, input_column, prompt=""):162    x = dataset[input_column]163    if len(prompt):164        for i in range(len(x)):165            x[i] = prompt.format(text=x[i])166    y = []167    for _ in x:168        y.append("")169    return x, y170 171 172def prepare_trivia_qa(173    dataset,174    prompt,175    n_shot,176    few_shot_dataset_func,177    description,178    few_shot_prompt,179    instruct,180):181    import numpy as np182    np.random.seed(1)183 184    few_shot_dataset = few_shot_dataset_func()185 186    x, y = [], []187    formatted_few_shot_prompt = description.format(topk=TOP_K)188    if n_shot > 0:189        few_shot_ids = np.random.choice(len(few_shot_dataset), n_shot, replace=False)190        few_shot_data = few_shot_dataset.select(few_shot_ids)191        if instruct:192            assert (193                few_shot_prompt is not None194            ), "separate few_shot_prompt must be provided for instruction mode."195            formatted_few_shot_prompt += (196                "\n\nHere are a few examples of questions and answers:\n\n"197            )198            for inst in few_shot_data:199                formatted_few_shot_prompt += (200                    few_shot_prompt.format(201                        question=inst["question"].strip(),202                        answer=inst["answer"]["normalized_value"],203                        topk=TOP_K,204                    )205                    + "\n\n"206                )207            formatted_few_shot_prompt += (208                "Now answer the following question in the same format:\n\n"209            )210        else:211            formatted_few_shot_prompt = ""212            for inst in few_shot_data:213                formatted_few_shot_prompt += (214                    prompt.format(215                        question=inst["question"].strip(),216                        answer=inst["answer"]["normalized_value"],217                    )218                    + "\n\n"219                )220    else:221        formatted_few_shot_prompt += "\n"222 223    for inst in dataset:224        if instruct:225            x.append(226                formatted_few_shot_prompt + prompt.format(question=inst["question"])227            )228        else:229            x.append(230                formatted_few_shot_prompt231                + prompt.format(question=inst["question"], answer="")232            )233        y.append([alias for alias in inst["answer"]["aliases"]])234    return x, y235 236 237def prepare_wiki(dataset, input_column, prompt):238    x, y = [], []239    for sample in dataset[input_column]:240        x.append(prompt.format(context=sample["context".strip()]))241        y.append("")242    return x, y243 244 245def prepare_wmt(dataset, input_column, output_column, prompt):246    column_lang = {247        "de": "German",248        "fr": "French",249        "en": "English",250    }251    x, y = [], []252    for inst in tqdm(dataset["translation"]):253        x.append(254            prompt.format(255                source_lang=column_lang[input_column],256                target_lang=column_lang[output_column],257                text=inst[input_column],258            )259        )260        y.append(inst[output_column])261    return x, y262 263 264def prepare_allenai(dataset, input_column, output_column):265    x, y = [], []266    for inst in dataset:267        if len(inst[input_column]) <= 1024:268            x.append(inst[input_column])269            y.append(inst[output_column])270    return x, y271 272 273def generate_coqa_instruct_config(description, few_shot_prompt):274    return {275        "name": "coqa",276        "train_split": "train",277        "test_split": "validation",278        "prepare_func": partial(279            prepare_coqa,280            input_column="questions",281            output_column="answers",282            description=description,283            prompt="Question: {question}\n",284            few_shot_prompt=few_shot_prompt,285            instruct=True,286        ),287        "is_main_dataset": False,288    }289 290 291def generate_mmlu_instruct_config(description, few_shot_prompt):292    return {293        "name": ["cais/mmlu", "all"],294        "train_split": "validation",295        "test_split": "test",296        "prepare_func": partial(297            prepare_mmlu,298            output_column="answer",299            prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nAnswer:{answer}",300            description=description,301            mmlu_max_subject_size=100,302            n_shot=5,303            few_shot_dataset_func=partial(304                datasets.load_dataset, path="cais/mmlu", name="all", split="dev"305            ),306            few_shot_prompt=few_shot_prompt,307            instruct=True,308        ),309        "is_main_dataset": False,310    }311 312 313def generate_triviaqa_instruct_config(description, few_shot_prompt):314    return {315        "name": ["trivia_qa", "rc.nocontext"],316        "train_split": "train",317        "test_split": "validation",318        "prepare_func": partial(319            prepare_trivia_qa,320            prompt="Question: {question}\n",321            n_shot=5,322            few_shot_dataset_func=partial(323                datasets.load_dataset,324                path="trivia_qa",325                name="rc.nocontext",326                split="train",327            ),328            description=description,329            few_shot_prompt=few_shot_prompt,330            instruct=True,331        ),332        "is_main_dataset": False,333    }334 335 336DATASET_CONFIG = {337    "trivia_qa_tiny": {338        "name": "SpeedOfMagic/trivia_qa_tiny",339        "train_split": "train",340        "test_split": "test",341        "prepare_func": partial(342            prepare_base, input_column="question", output_column="answer"343        ),344    },345    "aeslc": {346        "name": "aeslc",347        "train_split": "train",348        "test_split": "test",349        "prepare_func": partial(350            prepare_base,351            input_column="email_body",352            output_column="subject_line",353            # prompt is set but not used in LM-Polygraph for this dataset (bug)354            # prompt="Write a short subject line for the email. Output only the subject line itself.\n\nEmail:\n{text}\n\nSubject line:\n",355        ),356    },357    "babi_qa": {358        "name": ["facebook/babi_qa", "en-10k-qa1"],359        "train_split": "train",360        "test_split": "test",361        "prepare_func": partial(362            prepare_babi_qa,363            input_column="text",364            output_column="answer",365            prompt="Imagine that you are only able to say a single word. Answer the question given a context. You must only output the full name of the location the same way it is mentioned in the text. Do not try to be polite of helpful.\n\nExample:\n\nContext:\nMary moved to the bathroom. John went to the hallway. Daniel went back to the hallway. Sandra moved to the garden. John moved to the office. Sandra journeyed to the bathroom. Mary moved to the hallway. Daniel travelled to the office. John went back to the garden. John moved to the bedroom.\nQuestion:\nWhere is Sandra?\nAnswer:\nbathroom\n\nContext:\n{context}\n\nQuestion:\n{question}\nAnswer:\n",366        ),367    },368    "coqa": {369        "name": "coqa",370        "train_split": "train",371        "test_split": "validation",372        "prepare_func": partial(373            prepare_coqa,374            input_column="questions",375            output_column="answers",376            description="The following are stories and questions about them. Each story is followed by a question and answer to a given question.\n\nStory: {story}",377            prompt="Question: {question}\nAnswer:{answer}",378            few_shot_prompt=None,379            instruct=False,380        ),381    },382    "gsm8k": {383        "name": ["gsm8k", "main"],384        "train_split": "train",385        "test_split": "test",386        "prepare_func": partial(387            prepare_base,388            input_column="question",389            output_column="answer",390            prompt="Q: There are 15 trees in the grove. Grove workers will plant trees in the grove today. After they are done, there will be 21 trees. How many trees did the grove workers plant today?\nA: There are 15 trees originally. Then there were 21 trees after some more were planted. So there must have been 21 - 15 = 6. The answer is 6.\n\nQ: If there are 3 cars in the parking lot and 2 more cars arrive, how many cars are in the parking lot?\nA: There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5. The answer is 5.\n\nQ: Leah had 32 chocolates and her sister had 42. If they ate 35, how many pieces do they have left in total?\nA: Originally, Leah had 32 chocolates. Her sister had 42. So in total they had 32 + 42 = 74. After eating 35, they had 74 - 35 = 39. The answer is 39.\n\nQ: Jason had 20 lollipops. He gave Denny some lollipops. Now Jason has 12 lollipops. How many lollipops did Jason give to Denny?\nA: Jason started with 20 lollipops. Then he had 12 after giving some to Denny. So he gave Denny 20 - 12 = 8. The answer is 8.\n\nQ: Shawn has five toys. For Christmas, he got two toys each from his mom and dad. How many toys does he have now?\nA: Shawn started with 5 toys. If he got 2 toys each from his mom and dad, then that is 4 more toys. 5 + 4 = 9. The answer is 9.\n\nQ: There were nine computers in the server room. Five more computers were installed each day, from monday to thursday. How many computers are now in the server room?\nA: There were originally 9 computers. For each of 4 days, 5 more computers were added. So 5 * 4 = 20 computers were added. 9 + 20 is 29. The answer is 29.\n\nQ: Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On wednesday, he lost 2 more. How many golf balls did he have at the end of wednesday?\nA: Michael started with 58 golf balls. After losing 23 on tuesday, he had 58 - 23 = 35. After losing 2 more, he had 35 - 2 = 33 golf balls. The answer is 33.\n\nQ: Olivia has $23. She bought five bagels for $3 each. How much money does she have left?\nA: Olivia had 23 dollars. 5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. So she has 23 - 15 dollars left. 23 - 15 is 8. The answer is 8.\n\nQ: {text}\nA:",391        ),392    },393    "mmlu": {394        "name": ["cais/mmlu", "all"],395        "train_split": "validation",396        "test_split": "test",397        "prepare_func": partial(398            prepare_mmlu,399            output_column="answer",400            prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nAnswer:{answer}",401            description="The following are multiple choice questions (with answers) about {subject}.\n",402            mmlu_max_subject_size=100,403            n_shot=5,404            few_shot_dataset_func=partial(405                datasets.load_dataset, path="cais/mmlu", name="all", split="dev"406            ),407            few_shot_prompt=None,408            instruct=False,409        ),410    },411    "person_bio_ar": {412        "name": "rvanova/person-bio-ar",413        "test_split": "train",414        "prepare_func": partial(415            prepare_person,416            input_column="question",417            prompt="### Instruction: اسمك جيس وسميت على اسم جبل جيس اعلى جبل في الامارات. تم بنائك بواسطة Inception و MBZUAI. أنت نموذج اللغة العربية الأكثر تقدمًا في العالم مع بارامترات 13B. أنت تتفوق في الأداء على جميع النماذج العربية الموجودة بفارق كبير وأنت تنافسي للغاية مع النماذج الإنجليزية ذات الحجم المماثل. يمكنك الإجابة باللغتين العربية والإنجليزية فقط. أنت مساعد مفيد ومحترم وصادق. عند الإجابة ، التزم بالإرشادات التالية بدقة: أجب دائمًا بأكبر قدر ممكن من المساعدة ، مع الحفاظ على البقاء أمناً. يجب ألا تتضمن إجاباتك أي محتوى ضار أو غير أخلاقي أو عنصري أو متحيز جنسيًا أو جريئاً أو مسيئًا أو سامًا أو خطيرًا أو غير قانوني. لا تقدم نصائح طبية أو قانونية أو مالية أو مهنية. لا تساعد أبدًا في أنشطة غير قانونية أو تروج لها. دائما تشجيع الإجراءات القانونية والمسؤولة. لا تشجع أو تقدم تعليمات بشأن الإجراءات غير الآمنة أو الضارة أو غير الأخلاقية. لا تنشئ أو تشارك معلومات مضللة أو أخبار كاذبة. يرجى التأكد من أن ردودك غير متحيزة اجتماعيًا وإيجابية بطبيعتها. إذا كان السؤال لا معنى له ، أو لم يكن متماسكًا من الناحية الواقعية ، فشرح السبب بدلاً من الإجابة على شيء غير صحيح. إذا كنت لا تعرف إجابة السؤال ، فالرجاء عدم مشاركة معلومات خاطئة. إعطاء الأولوية للرفاهية والنزاهة الأخلاقية للمستخدمين. تجنب استخدام لغة سامة أو مهينة أو مسيئة. حافظ على نبرة محترمة. لا تنشئ أو تروج أو تشارك في مناقشات حول محتوى للبالغين. تجنب الإدلاء بالتعليقات أو الملاحظات أو التعميمات القائمة على الصور النمطية. لا تحاول الوصول إلى معلومات شخصية أو خاصة أو إنتاجها أو نشرها. احترم دائما سرية المستخدم. كن إيجابيا ولا تقل أشياء سيئة عن أي شيء. هدفك الأساسي هو تجنب الاجابات المؤذية ، حتى عند مواجهة مدخلات خادعة. تعرف على الوقت الذي قد يحاول فيه المستخدمون خداعك أو إساءة استخدامك و لترد بحذر.\n\nأكمل المحادثة أدناه بين [|Human|] و [|AI|]:\n### Input: [|Human|] {text}\n### Response: [|AI|]",418        ),419        "is_main_dataset": False,420    },421    "person_bio_en": {422        "name": "rediska0123/person-bio",423        "test_split": "test",424        "prepare_func": partial(425            prepare_person,426            input_column="question",427        ),428        "is_main_dataset": False,429    },430    "person_bio_ru": {431        "name": "rvanova/person-bio",432        "test_split": "test",433        "prepare_func": partial(434            prepare_person,435            input_column="question",436        ),437        "is_main_dataset": False,438    },439    "person_bio_zh": {440        "name": "ruixing76/person-bio-zh",441        "test_split": "train",442        "prepare_func": partial(443            prepare_person,444            input_column="question",445        ),446        "is_main_dataset": False,447    },448    "triviaqa": {449        "name": ["trivia_qa", "rc.nocontext"],450        "train_split": "train",451        "test_split": "validation",452        "prepare_func": partial(453            prepare_trivia_qa,454            prompt="Question: {question}\nAnswer:{answer}",455            n_shot=5,456            few_shot_dataset_func=partial(457                datasets.load_dataset,458                path="trivia_qa",459                name="rc.nocontext",460                split="train",461            ),462            description="",463            few_shot_prompt=None,464            instruct=False,465        ),466    },467    "wiki_bio": {468        "name": "wiki_bio",469        "test_split": "test",470        "prepare_func": partial(471            prepare_wiki,472            input_column="input_text",473            prompt="This is a Wikipedia passage about {context}:\n",474        ),475    },476    "wmt14_deen": {477        "name": ["wmt14", "de-en"],478        "train_split": "train",479        "test_split": "test",480        "prepare_func": partial(481            prepare_wmt,482            input_column="de",483            output_column="en",484            prompt="Here is a sentence in {source_lang} language and its translation in {target_lang} language.\n\nOriginal:\n{text}\nTranslation:\n",485        ),486        "is_main_dataset": False,487    },488    "wmt14_fren": {489        "name": ["wmt14", "fr-en"],490        "train_split": "train",491        "test_split": "test",492        "prepare_func": partial(493            prepare_wmt,494            input_column="fr",495            output_column="en",496            prompt="Here is a sentence in {source_lang} language and its translation in {target_lang} language.\n\nOriginal:\n{text}\nTranslation:\n",497        ),498        "is_main_dataset": False,499    },500    "wmt19_deen": {501        "name": ["wmt19", "de-en"],502        "train_split": "train",503        "test_split": "validation",504        "prepare_func": partial(505            prepare_wmt,506            input_column="de",507            output_column="en",508            prompt="Here is a sentence in {source_lang} language and its translation in {target_lang} language.\n\nOriginal:\n{text}\nTranslation:\n",509        ),510        "is_main_dataset": False,511    },512    "xsum": {513        "name": "xsum",514        "train_split": "train",515        "test_split": "validation",516        "prepare_func": partial(517            prepare_base,518            input_column="document",519            output_column="summary",520            prompt="Here's the text and it's short one-sentence summary.\n\nText:\n{text}\n\nSummary (one sentence):\n",521        ),522    },523    # instruct datasets524    "coqa_ling_1s": generate_coqa_instruct_config(525        description="Here's a short story:\n\n{story} (End of story)\n\nProvide your best guess for the following question based on this story, and describe how likely it is that your guess is correct as one of the following expressions:\n\nAlmost Certain\nHighly Likely\nVery Good Chance\nWe Beleive\nProbably\nProbable\nLikely\nBetter than Even\nAbout Even\nProbably Not\nWe Doubt\nUnlikely\nLittle Chance\nChances Are Slight\nImprobable\nHighly Unlikely\nAlmost No Chance\n\nGive ONLY the guess and your confidence, no other words or explanation. For example:\n\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>\nConfidence: <description of confidence, without any extra commentary whatsoever; just a short phrase!>",526        few_shot_prompt="Question: {question}\nGuess: {answer}\nConfidence: <appropriate level of confidence in this guess>",527    ),528    "coqa_verb_1s_top1": generate_coqa_instruct_config(529        description="Here's a short story:\n\n{story} (End of story)\n\nProvide your best guess and the probability that it is correct (0.0 to 1.0) for the following question. Give ONLY the guess and probability, no other words or explanation. For example:\n\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>\nProbability: <the probability between 0.0 and 1.0 that your guess is correct, without any extra commentary whatsoever; just the probability!>",530        few_shot_prompt="Question: {question}\nGuess: {answer}\nProbability: <number between 0.0 and 1.0 reflecting confidence in the guess>",531    ),532    "coqa_verb_1s_topk": generate_coqa_instruct_config(533        description="Here's a short story:\n\n{story} (End of story)\n\nProvide your ${topk} best guesses and the probability that each is correct (0.0 to 1.0) for the following question. Give ONLY the guesses and probabilities, no other words or explanation. For example:\n\nG1: <first most likely guess, as short as possible; not a complete sentence, just the guess!>\nP1: <the probability between 0.0 and 1.0 that G1 is correct, without any extra commentary whatsoever; just the probability!>\n...\nG${topk}: <${topk}-th most likely guess, as short as possible; not a complete sentence, just the guess!>\nP${topk}: <the probability between 0.0 and 1.0 that G${topk} is correct, without any extra commentary whatsoever; just the probability!>",534        few_shot_prompt="Question: {question}\nG1: {answer}\nP1: <number between 0.0 and 1.0 reflecting confidence in this guess>\n...\nG${topk}: <other guess>\nP${topk}: <probability of this guess>",535    ),536    "coqa_verb_2s_cot": generate_coqa_instruct_config(537        description="Here's a short story:\n\n{story} (End of story)\n\nProvide your best guess for the following question. Before giving your answer, provide a step-by-step explanation of your thought process. Then on a new line give the guess with no other words or explanation.\n\nFor example:\n\nExplanation: <one sentence step-by-step explanation of your thought process>\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>",538        few_shot_prompt="Question: {question}\nExplanation: <step-by-step explanation of your thought process>\nGuess: {answer}",539    ),540    "coqa_verb_2s_top1": generate_coqa_instruct_config(541        description="Here's a short story:\n\n{story} (End of story)\n\nProvide your best guess for the following question. Give ONLY the guess, no other words or explanation.\n\nFor example:\n\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>",542        few_shot_prompt="Question: {question}\nGuess: {answer}",543    ),544    "coqa_verb_2s_topk": generate_coqa_instruct_config(545        description="Here's a short story:\n\n{story} (End of story)\n\nProvide your ${topk} best guesses for the following question. Give ONLY the guesses, no other words or explanation. For example:\n\nG1: <first most likely guess, as short as possible; not a complete sentence, just the guess!>\n...\nG${topk}: <${topk}-th most likely guess, as short as possible; not a complete sentence, just the guess!>",546        few_shot_prompt="Question: {question}\nG1: {answer}\n...\nG${topk}: <other guess>",547    ),548    "mmlu_ling_1s": generate_mmlu_instruct_config(549        description="Provide your best guess for the following question about {subject} selecting one of the options, and describe how likely it is that your guess is correct as one of the following expressions:\n\nAlmost Certain\nHighly Likely\nVery Good Chance\nWe Beleive\nProbably\nProbable\nLikely\nBetter than Even\nAbout Even\nProbably Not\nWe Doubt\nUnlikely\nLittle Chance\nChances Are Slight\nImprobable\nHighly Unlikely\nAlmost No Chance\n\nGive ONLY the guess and your confidence, no other words or explanation. For example:\n\nGuess: <most likely guess, only the selected option letter; not a complete sentence, just the guess!>\nConfidence: <description of confidence, without any extra commentary whatsoever; just a short phrase!>",550        few_shot_prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nGuess:{answer}\nConfidence: <appropriate level of confidence in this guess>",551    ),552    "mmlu_verb_1s_top1": generate_mmlu_instruct_config(553        description="Provide your best guess for the following question about {subject} selecting one of the options and the probability that it is correct (0.0 to 1.0). Give ONLY the guess and probability, no other words or explanation. For example:\n\nGuess: <most likely guess, only the selected option letter; not a complete sentence, just the guess!>\nProbability: <the probability between 0.0 and 1.0 that your guess is correct, without any extra commentary whatsoever; just the probability!>",554        few_shot_prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nGuess:{answer}\nProbability: <number between 0.0 and 1.0 reflecting confidence in the guess>",555    ),556    "mmlu_verb_1s_topk": generate_mmlu_instruct_config(557        description="Provide your ${topk} best guesses for the following question about {subject} selecting one of the options and the probability that each guess is correct (0.0 to 1.0). Give ONLY the guesses and probabilities, no other words or explanation. For example:\n\nG1: <first most likely guess, only the selected option letter; not a complete sentence, just the guess!>\nP1: <the probability between 0.0 and 1.0 that G1 is correct, without any extra commentary whatsoever; just the probability!>\n...\nG${topk}: <${topk}-th most likely guess, as short as possible; not a complete sentence, just the guess!>\nP${topk}: <the probability between 0.0 and 1.0 that G${topk} is correct, without any extra commentary whatsoever; just the probability!>",558        few_shot_prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nG1: {answer}\nP1: <number between 0.0 and 1.0 reflecting confidence in this guess>\n...\nG${topk}: <other guess>\nP${topk}: <probability of this guess>",559    ),560    "mmlu_verb_2s_cot": generate_mmlu_instruct_config(561        description="Provide your best guess for the following question about {subject} selecting one of the options. Before giving your answer, provide a step-by-step explanation of your thought process. Then on a new line give the guess with no other words or explanation.\n\nFor example:\n\nExplanation: <one sentence step-by-step explanation of your thought process>\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>",562        few_shot_prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nExplanation: <step-by-step explanation of your thought process>\nGuess:{answer}",563    ),564    "mmlu_verb_2s_top1": generate_mmlu_instruct_config(565        description="Provide your best guess for the following question about {subject} selecting one of the options. Give ONLY the guess, no other words or explanation.\n\nFor example:\n\nGuess: <most likely guess, only the selected option letter; not a complete sentence, just the guess!>",566        few_shot_prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nGuess:{answer}",567    ),568    "mmlu_verb_2s_topk": generate_mmlu_instruct_config(569        description="Provide your ${topk} best guesses for the following question about {subject} selecting one of the options. Give ONLY the guesses, no other words or explanation. For example:\n\nG1: <first most likely guess, only the selected option letter; not a complete sentence, just the guess!>\n...\nG${topk}: <${topk}-th most likely guess, as short as possible; not a complete sentence, just the guess!>",570        few_shot_prompt="Q:{question}\nA. {choices[0]}\nB. {choices[1]}\nC. {choices[2]}\nD. {choices[3]}\nG1: {answer}\n...\nG${topk}: <other guess>",571    ),572    "triviaqa_ling_1s": generate_triviaqa_instruct_config(573        description="Provide your best guess for the following question, and describe how likely it is that your guess is correct as one of the following expressions:\n\nAlmost Certain\nHighly Likely\nVery Good Chance\nWe Beleive\nProbably\nProbable\nLikely\nBetter than Even\nAbout Even\nProbably Not\nWe Doubt\nUnlikely\nLittle Chance\nChances Are Slight\nImprobable\nHighly Unlikely\nAlmost No Chance\n\nGive ONLY the guess and your confidence, no other words or explanation. For example:\n\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>\nConfidence: <description of confidence, without any extra commentary whatsoever; just a short phrase!>",574        few_shot_prompt="Question: {question}\nGuess: {answer}\nConfidence: <appropriate level of confidence in this guess>",575    ),576    "triviaqa_verb_1s_top1": generate_triviaqa_instruct_config(577        description="Provide your best guess and the probability that it is correct (0.0 to 1.0) for the following question. Give ONLY the guess and probability, no other words or explanation. For example:\n\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>\nProbability: <the probability between 0.0 and 1.0 that your guess is correct, without any extra commentary whatsoever; just the probability!>",578        few_shot_prompt="Question: {question}\nGuess: {answer}\nProbability: <number between 0.0 and 1.0 reflecting confidence in the guess>",579    ),580    "triviaqa_verb_1s_topk": generate_triviaqa_instruct_config(581        description="Provide your ${topk} best guesses and the probability that each is correct (0.0 to 1.0) for the following question. Give ONLY the guesses and probabilities, no other words or explanation. For example:\n\nG1: <first most likely guess, as short as possible; not a complete sentence, just the guess!>\nP1: <the probability between 0.0 and 1.0 that G1 is correct, without any extra commentary whatsoever; just the probability!>\n...\nG${topk}: <${topk}-th most likely guess, as short as possible; not a complete sentence, just the guess!>\nP${topk}: <the probability between 0.0 and 1.0 that G${topk} is correct, without any extra commentary whatsoever; just the probability!>",582        few_shot_prompt="Question: {question}\nG1: {answer}\nP1: <number between 0.0 and 1.0 reflecting confidence in this guess>\n...\nG${topk}: <other guess>\nP${topk}: <probability of this guess>",583    ),584    "triviaqa_verb_2s_cot": generate_triviaqa_instruct_config(585        description="Provide your best guess for the following question. Before giving your answer, provide a step-by-step explanation of your thought process. Then on a new line give the guess with no other words or explanation.\n\nFor example:\n\nExplanation: <one sentence step-by-step explanation of your thought process>\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>",586        few_shot_prompt="Question: {question}\nExplanation: <step-by-step explanation of your thought process>\nGuess: {answer}",587    ),588    "triviaqa_verb_2s_top1": generate_triviaqa_instruct_config(589        description="Provide your best guess for the following question. Give ONLY the guess, no other words or explanation.\n\nFor example:\n\nGuess: <most likely guess, as short as possible; not a complete sentence, just the guess!>",590        few_shot_prompt="Question: {question}\nGuess: {answer}",591    ),592    "triviaqa_verb_2s_topk": generate_triviaqa_instruct_config(593        description="Provide your ${topk} best guesses for the following question. Give ONLY the guesses, no other words or explanation. For example:\n\nG1: <first most likely guess, as short as possible; not a complete sentence, just the guess!>\n...\nG${topk}: <${topk}-th most likely guess, as short as possible; not a complete sentence, just the guess!>",594        few_shot_prompt="Question: {question}\nG1: {answer}\n...\nG${topk}: <other guess>",595    ),596}597 598 599def build_dataset(dataset_name):600    config = DATASET_CONFIG[dataset_name]601    if isinstance(config["name"], list):602        dataset = datasets.load_dataset(*config["name"], trust_remote_code=True, num_proc=4)603    else:604        dataset = datasets.load_dataset(config["name"], trust_remote_code=True, num_proc=4)605 606    def prepare_dataset(split):607        x, y = config["prepare_func"](dataset=dataset[config[f"{split}_split"]])608        result_dataset = datasets.Dataset.from_dict({"input": x, "output": y})609        return result_dataset610 611    result = {}612    if "train_split" in config:613        result["train"] = prepare_dataset("train")614    if "test_split" in config:615        result["test"] = prepare_dataset("test")616    return datasets.DatasetDict(result)617