"""Preprocessing functions for various benchmark datasets. This module provides data loading and prompt formatting functions for: - Math benchmarks: MATH, GSM8K, AIME, Minerva Math, OmniMath, etc. - Coding benchmarks: HumanEval, LiveCodeBench, MBPP - Multiple-choice: MMLU, MMLU Pro, GPQA - Instruction following: IFEval, IFBench, MT-Bench - General: AlpacaEval, Arena-Hard """ import json import pandas import os def preprocess_gpqa_chatml_template(data_file, use_r1=False, think=True): """Preprocess GPQA dataset with ChatML template formatting. Args: data_file: Path to GPQA JSON file use_r1: Whether to use DeepSeek R1-style prompting (default: False) think: Whether to enable thinking mode (default: True) Returns: list: Formatted prompts with ChatML template """ if use_r1: QUERY_TEMPLATE_MULTICHOICE = "{Question}\n\n\nA. {choice1}\nB. {choice2}\nC. {choice3}\nD. {choice4}\n\nPlease reason step-by-step and put your choice letter without any other text with \\boxed{{}} in the end. Let's think step by step and output the final answer within \\boxed{{}}." else: QUERY_TEMPLATE_MULTICHOICE = "Return your final response within \\boxed{{}} and only include the letter choice (e.g., A, B, C, or D) as your final response.\n\n{Question}\n\nAnswer Choices:\n(A) {choice1}\n(B) {choice2}\n(C) {choice3}\n(D) {choice4}" instruction = "<|im_start|>system\nYou are a helpful and harmless assistant.<|im_end|>\n" with open(data_file, "r") as f: data_list = json.load(f) prompt_list = [] for item in data_list: choices_dict = dict( Question=item['question'].strip(), choice1=item['choice_A'].strip(), choice2=item['choice_B'].strip(), choice3=item['choice_C'].strip(), choice4=item['choice_D'].strip() ) final_question = QUERY_TEMPLATE_MULTICHOICE.format(**choices_dict) if use_r1: final_prompt = """<｜begin▁of▁sentence｜><｜User｜>{question}.<｜Assistant｜>\n""".format(question=final_question) else: if think: final_prompt = instruction + "<|im_start|>user\n" + final_question + " /think<|im_end|>\n<|im_start|>assistant\n\n" else: final_prompt = instruction + "<|im_start|>user\n" + final_question + " /no_think<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list def preprocess_gpqa_raw_template(data_file, use_r1=False, think=True): """Preprocess GPQA dataset with raw (no template) formatting. Args: data_file: Path to GPQA JSON file use_r1: Whether to use DeepSeek R1-style prompting (default: False) think: Whether to enable thinking mode (default: True) Returns: list: Raw formatted prompts without chat template """ if use_r1: QUERY_TEMPLATE_MULTICHOICE = "{Question}\n\n\nA. {choice1}\nB. {choice2}\nC. {choice3}\nD. {choice4}\n\nPlease reason step-by-step and put your choice letter without any other text with \\boxed{{}} in the end. Let's think step by step and output the final answer within \\boxed{{}}." else: QUERY_TEMPLATE_MULTICHOICE = "Return your final response within \\boxed{{}} and only include the letter choice (e.g., A, B, C, or D) as your final response.\n\n{Question}\n\nAnswer Choices:\n(A) {choice1}\n(B) {choice2}\n(C) {choice3}\n(D) {choice4}" with open(data_file, "r") as f: data_list = json.load(f) prompt_list = [] for item in data_list: choices_dict = dict( Question=item['question'].strip(), choice1=item['choice_A'].strip(), choice2=item['choice_B'].strip(), choice3=item['choice_C'].strip(), choice4=item['choice_D'].strip() ) final_question = QUERY_TEMPLATE_MULTICHOICE.format(**choices_dict) prompt_list.append(final_question) return prompt_list def preprocess_gsm8k_zeroshot_chatml_template(data_file): """Preprocess GSM8K dataset with zero-shot ChatML template. Args: data_file: Path to GSM8K JSON file Returns: list: Formatted prompts with ChatML template and thinking enabled """ with open(data_file, "r") as f: data_list = json.load(f) instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] for item in data_list: final_question = item['question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_gsm8k_zeroshot_raw(data_file): """Preprocess GSM8K dataset with zero-shot raw formatting. Args: data_file: Path to GSM8K JSON file Returns: list: Raw question prompts without chat template """ with open(data_file, "r") as f: data_list = json.load(f) prompt_list = [] for item in data_list: final_question = item['question'].strip() final_prompt = final_question prompt_list.append(final_prompt) return prompt_list def preprocess_humaneval_raw(data_file): """Preprocess HumanEval code generation dataset. Args: data_file: Path to HumanEval JSON file Returns: tuple: (prompt_list, qid_list) - prompt_list: Code completion prompts with instructions - qid_list: Task IDs """ qid_list = [] prompt_list = [] instruction = "Read the following function signature and docstring, and fully implement the function described. Your response should only contain the code for this function.\n" with open(data_file, "r") as f: data_dict = json.load(f) for key, values in data_dict.items(): qid_list.append(key) prompt = instruction + values['prompt'] prompt_list.append(prompt) return prompt_list, qid_list def preprocess_math_zeroshot_chatml_template(data_file): """Preprocess MATH dataset with zero-shot ChatML template. Args: data_file: Path to MATH CSV file Returns: list: Formatted prompts with ChatML template and thinking enabled """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" df = pandas.read_csv(data_file) test_list = [row.to_dict() for _, row in df.iterrows()] prompt_list = [] for item in test_list: final_question = item['Question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_math500_zeroshot_chatml_template(data_file, use_r1=False): """Preprocess MATH-500 dataset with zero-shot prompting. Args: data_file: Path to MATH-500 JSONL file use_r1: Whether to use DeepSeek R1-style prompting (default: False) Returns: list: Formatted prompts with boxed answer instruction """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: data_dict = json.loads(line) final_question = data_dict['problem'].strip() if use_r1: final_prompt = """<｜begin▁of▁sentence｜><｜User｜>{question}\nPlease reason step by step, and put your final answer within \boxed{{}}.<｜Assistant｜>\n""".format(question=final_question) else: final_prompt = instruction + "<|im_start|>user\n" + final_question + "\n\nPlease place your final answer inside \\boxed{}." + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_minerva_math_chatml_template(data_file): """Preprocess Minerva Math dataset with ChatML template. Args: data_file: Path to Minerva Math JSONL file Returns: list: Formatted prompts with boxed answer instruction """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['problem'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "\n\nPlease place your final answer inside \\boxed{}." + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_gaokao2023en_chatml_template(data_file): """Preprocess Gaokao 2023 English dataset with ChatML template. Args: data_file: Path to Gaokao 2023 English JSONL file Returns: list: Formatted prompts with boxed answer instruction """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r", encoding="utf-8") as f: for line in f: item = json.loads(line) final_question = item['question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "\n\nPlease place your final answer inside \\boxed{}." + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_olympiadbench_chatml_template(data_file): """Preprocess OlympiadBench dataset with ChatML template. Args: data_file: Path to OlympiadBench JSONL file Returns: list: Formatted prompts with boxed answer instruction """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r", encoding="utf-8") as f: for line in f: item = json.loads(line) final_question = item['question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "\n\nPlease place your final answer inside \\boxed{}." + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_collegemath_chatml_template(data_file): """Preprocess College Math dataset with ChatML template. Args: data_file: Path to College Math JSONL file Returns: list: Formatted prompts with boxed answer instruction """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "\n\nPlease place your final answer inside \\boxed{}." + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_aime24_chatml_template(data_file): """Preprocess AIME 2024 dataset with ChatML template. Args: data_file: Path to AIME 2024 JSONL file Returns: list: Formatted prompts with thinking enabled """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_aime25_chatml_template(data_file): """Preprocess AIME 2025 dataset with ChatML template. Args: data_file: Path to AIME 2025 JSONL file Returns: list: Formatted prompts with thinking enabled """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['problem'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_aime24_raw(data_file): """Preprocess AIME 2024 dataset with raw formatting. Args: data_file: Path to AIME 2024 JSONL file Returns: list: Raw prompts with boxed answer instruction """ prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['question'].strip() final_prompt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}.".format(question=final_question) prompt_list.append(final_prompt) return prompt_list def preprocess_aime25_raw(data_file): """Preprocess AIME 2025 dataset with raw formatting. Args: data_file: Path to AIME 2025 JSONL file Returns: list: Raw prompts with boxed answer instruction """ prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['problem'].strip() final_prompt = "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}.".format(question=final_question) prompt_list.append(final_prompt) return prompt_list def preprocess_amc23_chatml_template(data_file): """Preprocess AMC 2023 dataset with ChatML template. Args: data_file: Path to AMC 2023 JSONL file Returns: list: Formatted prompts with thinking enabled """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['question'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_omnimath_chatml_template(data_file): """Preprocess OmniMath dataset with ChatML template. Args: data_file: Path to OmniMath JSONL file Returns: list: Formatted prompts with thinking enabled """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" prompt_list = [] with open(data_file, "r") as f: for line in f: item = json.loads(line) final_question = item['problem'].strip() final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list def preprocess_ifeval_chatml_template(data_file): """Preprocess IFEval instruction-following dataset with ChatML template. Args: data_file: Path to IFEval JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: Formatted prompts with ChatML template - qid_list: Task IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['key']) first_question = item['prompt'] final_prompt = "<|im_start|>user\n" + first_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_ifeval_raw(data_file): """Preprocess IFEval instruction-following dataset with raw formatting. Args: data_file: Path to IFEval JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: Raw instruction prompts - qid_list: Task IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['key']) final_prompt = item['prompt'] prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_ifbench_raw(data_file): """Preprocess IFBench instruction-following dataset with raw formatting. Args: data_file: Path to IFBench JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: Raw instruction prompts - qid_list: Task IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['key']) final_prompt = item['prompt'] prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_arena_hard_chatml_template(data_file): """Preprocess Arena-Hard dataset with ChatML template. Args: data_file: Path to Arena-Hard JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: Formatted prompts with ChatML template - qid_list: Question IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['question_id']) first_question = item['turns'][0]['content'] final_prompt = "<|im_start|>user\n" + first_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_arena_hard_raw(data_file): """Preprocess Arena-Hard dataset with raw formatting. Args: data_file: Path to Arena-Hard JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: Raw question prompts - qid_list: Question IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['question_id']) first_question = item['turns'][0]['content'] prompt_list.append(first_question) return prompt_list, qid_list def preprocess_arena_hard_v2_raw(data_file): """Preprocess Arena-Hard v2.0 dataset with raw formatting. Args: data_file: Path to Arena-Hard v2.0 JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: Raw question prompts - qid_list: Unique IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['uid']) first_question = item['prompt'] prompt_list.append(first_question) return prompt_list, qid_list def preprocess_alpaca_eval_raw(data_file): """Preprocess AlpacaEval dataset with raw formatting. Args: data_file: Path to AlpacaEval JSON file Returns: tuple: (prompt_list, qid_list) - prompt_list: Raw instruction prompts - qid_list: Sequential question IDs """ with open(data_file, "r") as f: data_list = json.load(f) qid_list = [] prompt_list = [] for i, item in enumerate(data_list): qid_list.append(i) final_prompt = item['instruction'] prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_alpaca_eval_chatml_template(data_file): """Preprocess AlpacaEval dataset with ChatML template. Args: data_file: Path to AlpacaEval JSON file Returns: tuple: (prompt_list, qid_list) - prompt_list: Formatted prompts with ChatML template - qid_list: Sequential question IDs """ with open(data_file, "r") as f: data_list = json.load(f) qid_list = [] prompt_list = [] for i, item in enumerate(data_list): qid_list.append(i) first_question = item['instruction'] final_prompt = "<|im_start|>user\n" + first_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_mtbench_firstturn(data_file): """Preprocess MT-Bench first turn with ChatML template. Args: data_file: Path to MT-Bench JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: First turn prompts with ChatML template - qid_list: Question IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['question_id']) first_question = item['turns'][0] final_prompt = "<|im_start|>user\n" + first_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_mtbench_firstturn_raw(data_file): """Preprocess MT-Bench first turn with raw formatting. Args: data_file: Path to MT-Bench JSONL file Returns: tuple: (prompt_list, qid_list) - prompt_list: First turn raw prompts - qid_list: Question IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['question_id']) first_question = item['turns'][0] prompt_list.append(first_question) return prompt_list, qid_list def preprocess_mtbench_secondturn(data_file, output_file): """Preprocess MT-Bench second turn with ChatML template. Args: data_file: Path to MT-Bench JSONL file output_file: Model output file for the first turn of MT-Bench Returns: tuple: (prompt_list, qid_list) - prompt_list: Second turn prompts with conversation history - qid_list: Question IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] id2output = {} with open(output_file, "r") as f: for line in f: item = json.loads(line) id2output[item['task_id']] = item['output'] qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['question_id']) first_question = item['turns'][0] second_question = item['turns'][1] model_output = id2output[item['question_id']] final_prompt = "<|im_start|>user\n" + first_question + "<|im_end|>\n<|im_start|>assistant\n" + model_output + "<|im_end|>\n" + \ "<|im_start|>user\n" + second_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_mtbench_secondturn_raw(data_file, output_file): """Preprocess MT-Bench second turn with raw formatting. Args: data_file: Path to MT-Bench JSONL file output_file: Model output file for the first turn of MT-Bench Returns: tuple: (prompt_list, qid_list) - prompt_list: Second turn prompts as chat message lists - qid_list: Question IDs """ with open(data_file, "r") as f: data = f.readlines() data_list = [json.loads(x) for x in data] id2output = {} with open(output_file, "r") as f: for line in f: item = json.loads(line) id2output[item['task_id']] = (item['output'], item['reason_text']) qid_list = [] prompt_list = [] for item in data_list: qid_list.append(item['question_id']) first_question = item['turns'][0] second_question = item['turns'][1] output, reason = id2output[item['question_id']] model_output = output chat = [ {'role': 'user', 'content': first_question}, {'role': 'assistant', 'content': model_output}, {'role': 'user', 'content': second_question} ] prompt_list.append(chat) return prompt_list, qid_list def preprocess_mmlu_chatml_template(data_file): """Preprocess MMLU dataset with few-shot ChatML template. Args: data_file: Path to MMLU CSV file Returns: list: Formatted prompts with few-shot examples and ChatML template """ def _load_mmlu_cot_fewshot_examples(): import yaml current_folder = os.path.dirname(os.path.abspath(__file__)) fewshot_folder = os.path.join(current_folder, "flan_cot_fewshot") file_list = os.listdir(fewshot_folder) fewshot_dict = {} for filename in file_list: with open(os.path.join(fewshot_folder, filename)) as f: data = yaml.safe_load(f) dataset_name = data['dataset_name'].strip() description = data['description'].strip() sample_list = data['fewshot_config']["samples"] prompt = description for sample in sample_list: prompt += "\n\n" prompt += "Q: " + sample['question'].strip() + "\n" + "A: " + sample['target'].strip() fewshot_dict[dataset_name] = prompt return fewshot_dict fewshot_dict = _load_mmlu_cot_fewshot_examples() df = pandas.read_csv(data_file) test_list = [row.to_dict() for _, row in df.iterrows()] prompt_list = [] for item in test_list: subject = item['Subject'] fewshot_prompt = fewshot_dict[subject] question = item['Question'] choice_a = str(item['A']).strip() choice_b = str(item['B']).strip() choice_c = str(item['C']).strip() choice_d = str(item['D']).strip() final_question = fewshot_prompt + "\n\n" + "Q: " + question + "\n" final_question += "(A) " + choice_a + " (B) " + choice_b + " (C) " + choice_c + " (D) " + choice_d + "\n" final_question += "A: " final_prompt = "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list def preprocess_mmlu_raw_template(data_file): """Preprocess MMLU dataset with few-shot raw formatting. Args: data_file: Path to MMLU CSV file Returns: list: Raw prompts with few-shot examples """ def _load_mmlu_cot_fewshot_examples(): import yaml current_folder = os.path.dirname(os.path.abspath(__file__)) fewshot_folder = os.path.join(current_folder, "flan_cot_fewshot") file_list = os.listdir(fewshot_folder) fewshot_dict = {} for filename in file_list: with open(os.path.join(fewshot_folder, filename)) as f: data = yaml.safe_load(f) dataset_name = data['dataset_name'].strip() description = data['description'].strip() sample_list = data['fewshot_config']["samples"] prompt = description for sample in sample_list: prompt += "\n\n" prompt += "Q: " + sample['question'].strip() + "\n" + "A: " + sample['target'].strip() fewshot_dict[dataset_name] = prompt return fewshot_dict fewshot_dict = _load_mmlu_cot_fewshot_examples() df = pandas.read_csv(data_file) test_list = [row.to_dict() for _, row in df.iterrows()] prompt_list = [] for item in test_list: subject = item['Subject'] fewshot_prompt = fewshot_dict[subject] question = item['Question'] choice_a = str(item['A']).strip() choice_b = str(item['B']).strip() choice_c = str(item['C']).strip() choice_d = str(item['D']).strip() final_question = fewshot_prompt + "\n\n" + "Q: " + question + "\n" final_question += "(A) " + choice_a + " (B) " + choice_b + " (C) " + choice_c + " (D) " + choice_d + "\n" final_question += "A: " prompt_list.append(final_question) return prompt_list QUERY_TEMPLATE_MULTICHOICE = """ Answer the following multiple choice question. The last line of your response should be of the following format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before answering. {Question} A) {A} B) {B} C) {C} D) {D} """.strip() def preprocess_mmlu_r1_raw_template(data_file): """Preprocess MMLU dataset for R1-style models with zero-shot prompting. Args: data_file: Path to MMLU CSV file Returns: list: Formatted prompts for R1-style reasoning models """ df = pandas.read_csv(data_file) test_list = [row.to_dict() for _, row in df.iterrows()] prompt_list = [] for item in test_list: subject = item['Subject'] question = item['Question'] choice_a = str(item['A']).strip() choice_b = str(item['B']).strip() choice_c = str(item['C']).strip() choice_d = str(item['D']).strip() # final_question = fewshot_prompt + "\n\n" + "Q: " + question + "\n" # final_question += "(A) " + choice_a + " (B) " + choice_b + " (C) " + choice_c + " (D) " + choice_d + "\n" # final_question += "A: " final_question = QUERY_TEMPLATE_MULTICHOICE.format(Question=question, A=choice_a, B=choice_b, C=choice_c, D=choice_d) prompt_list.append(final_question) return prompt_list def preprocess_mmlu_r1_raw_template_wdai(data_file): """Preprocess MMLU dataset with boxed answer format (alternative R1 style). Args: data_file: Path to MMLU CSV file Returns: list: Formatted prompts with boxed answer instruction """ MMLU_QUERY_TEMPLATE_MULTICHOICE = "Answer the following multiple-choice question. At the end of your response, conclude with the sentence `The answer is \\boxed{{X}}.`, replacing X with the correct capital letter of your choice.\n\n{Question}\n\nAnswer Choices:\n(A) {choice1}\n(B) {choice2}\n(C) {choice3}\n(D) {choice4}" df = pandas.read_csv(data_file) test_list = [row.to_dict() for _, row in df.iterrows()] prompt_list = [] for item in test_list: choices_dict = dict( Question=item['Question'].strip(), choice1=str(item['A']).strip(), choice2=str(item['B']).strip(), choice3=str(item['C']).strip(), choice4=str(item['D']).strip() ) final_question = MMLU_QUERY_TEMPLATE_MULTICHOICE.format(**choices_dict) prompt_list.append(final_question) return prompt_list def preprocess_mmlu_pro_chatml_template(data_file, fewshot_file): """Preprocess MMLU-Pro dataset with 5-shot ChatML template. Args: data_file: Path to MMLU-Pro test JSON file fewshot_file: Path to MMLU-Pro validation JSON file (for few-shot examples) Returns: list: Formatted prompts with 5-shot examples and ChatML template """ def _preprocess(data_list): output_list = [] for item in data_list: options = [] for opt in item["options"]: if opt == "N/A": continue options.append(opt) item["options"] = options output_list.append(item) return output_list def _categorize_basedon_subject(data_list): fewshot_dict = {} for item in data_list: subject = item['category'] if subject in fewshot_dict: fewshot_dict[subject].append(item) else: fewshot_dict[subject] = [item] return fewshot_dict def _format_each_sample(sample, is_test): choices = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P"] question = sample['question'] options = sample['options'] sample_prompt = "Q: " + question + "\n" for i, opt in enumerate(options): sample_prompt += "(%s) %s " % (choices[i], opt) sample_prompt = sample_prompt.strip() + "\n" if is_test: sample_prompt += "A: " else: sample_prompt += sample['cot_content'].strip() return sample_prompt def _get_fewshot_prompt(fewshot_samples, test_sample, subject): description = "The following are multiple choice questions (with answers) about %s." % subject final_question = description + "\n\n" for sample in fewshot_samples: sample_prompt = _format_each_sample(sample, is_test=False) final_question += sample_prompt + "\n\n" test_prompt = _format_each_sample(test_sample, is_test=True) final_question += test_prompt return final_question with open(fewshot_file, "r") as f: fewshot_list = json.load(f) with open(data_file, "r") as f: test_list = json.load(f) fewshot_list = _preprocess(fewshot_list) fewshot_dict = _categorize_basedon_subject(fewshot_list) test_list = _preprocess(test_list) prompt_list = [] for test_sample in test_list: subject = test_sample['category'] fewshot_samples = fewshot_dict[subject] ## 5-shot examples assert len(fewshot_samples) == 5 final_question = _get_fewshot_prompt(fewshot_samples, test_sample, subject) final_prompt = "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list def preprocess_mmlu_pro_zero_shot_chatml_template(data_file, think=True): """Preprocess MMLU-Pro dataset with zero-shot ChatML template. Args: data_file: Path to MMLU-Pro test JSON file think: Whether to enable thinking mode (default: True) Returns: list: Formatted prompts with ChatML template and boxed answer instruction """ def _preprocess(data_list): output_list = [] for item in data_list: options = [] for opt in item["options"]: if opt == "N/A": continue options.append(opt) item["options"] = options output_list.append(item) return output_list def _format_each_sample(sample): choices = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P"] question = sample['question'] options = sample['options'] # sample_prompt = "Answer the following multiple-choice question. At the end of your response, conclude with the sentence `The answer is \\boxed{{X}}.`, replacing X with the correct capital letter of your choice.\n\n" # sample_prompt += question + "\n\nAnswer Choices:" # for i, opt in enumerate(options): # sample_prompt += "\n(%s) %s" % (choices[i], opt) # sample_prompt = sample_prompt.strip() + "\n" sample_prompt = "Question:\n" + question + "\n\nAnswer Choices:" for i, opt in enumerate(options): sample_prompt += "\n(%s) %s" % (choices[i], opt) sample_prompt += "\n\nConclude your response with the sentence `The answer is \\boxed{{X}}.`, in which X is the correct capital letter of your choice." sample_prompt = sample_prompt.strip() + "\n" return sample_prompt instruction = "<|im_start|>system\nYou are a helpful and harmless assistant.<|im_end|>\n" with open(data_file, "r") as f: test_list = json.load(f) test_list = _preprocess(test_list) prompt_list = [] for test_sample in test_list: test_prompt = _format_each_sample(test_sample) if think: final_prompt = instruction + "<|im_start|>user\n" + test_prompt + "\n /think<|im_end|>\n<|im_start|>assistant\n\n" else: final_prompt = instruction + "<|im_start|>user\n" + test_prompt + "\n /no_think<|im_end|>\n<|im_start|>assistant\n" prompt_list.append(final_prompt) return prompt_list def preprocess_mmlu_pro_zero_shot_raw_template(data_file, think=True): """Preprocess MMLU-Pro dataset with zero-shot raw formatting. Args: data_file: Path to MMLU-Pro test JSON file think: Whether to enable thinking mode (default: True, currently unused) Returns: list: Raw prompts with boxed answer instruction """ def _preprocess(data_list): output_list = [] for item in data_list: options = [] for opt in item["options"]: if opt == "N/A": continue options.append(opt) item["options"] = options output_list.append(item) return output_list def _format_each_sample(sample): choices = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P"] question = sample['question'] options = sample['options'] # sample_prompt = "Answer the following multiple-choice question. At the end of your response, conclude with the sentence `The answer is \\boxed{{X}}.`, replacing X with the correct capital letter of your choice.\n\n" # sample_prompt += question + "\n\nAnswer Choices:" # for i, opt in enumerate(options): # sample_prompt += "\n(%s) %s" % (choices[i], opt) # sample_prompt = sample_prompt.strip() + "\n" sample_prompt = "Question:\n" + question + "\n\nAnswer Choices:" for i, opt in enumerate(options): sample_prompt += "\n(%s) %s" % (choices[i], opt) sample_prompt += "\n\nConclude your response with the sentence `The answer is \\boxed{{X}}.`, in which X is the correct capital letter of your choice." sample_prompt = sample_prompt.strip() + "\n" return sample_prompt with open(data_file, "r") as f: test_list = json.load(f) test_list = _preprocess(test_list) prompt_list = [] for test_sample in test_list: test_prompt = _format_each_sample(test_sample) prompt_list.append(test_prompt) return prompt_list def preprocess_livecodebench_chatml_template(data_file): """Preprocess LiveCodeBench dataset with ChatML template. Args: data_file: Path to LiveCodeBench JSON file Returns: tuple: (prompt_list, qid_list) - prompt_list: Formatted coding prompts with ChatML template - qid_list: Question IDs """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" code_instruction_nostartercode = """Write Python code to solve the problem. Please place the solution code in the following format:\n```python\n# Your solution code here\n```""" code_instruction_hasstartercode = """Please place the solution code in the following format:\n```python\n# Your solution code here\n```""" with open(data_file, "r") as f: data_list = json.load(f) prompt_list = [] qid_list = [] for item in data_list: question = item['question_content'].strip() if item['starter_code'] != "": question += "\n\n" + "Solve the problem starting with the provided function header.\n\nFunction header:\n" + "```\n" + item['starter_code'] + "\n```" question += "\n\n" + code_instruction_hasstartercode else: question += "\n\n" + code_instruction_nostartercode final_prompt = instruction + "<|im_start|>user\n" + question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) qid_list.append(item['question_id']) return prompt_list, qid_list def preprocess_livecodebench_raw(data_file): """Preprocess LiveCodeBench dataset with raw formatting. Args: data_file: Path to LiveCodeBench JSON file Returns: tuple: (prompt_list, qid_list) - prompt_list: Raw coding prompts - qid_list: Question IDs """ code_instruction_nostartercode = """Write Python code to solve the problem. Please place the solution code in the following format:\n```python\n# Your solution code here\n```""" code_instruction_hasstartercode = """Please place the solution code in the following format:\n```python\n# Your solution code here\n```""" with open(data_file, "r") as f: data_list = json.load(f) prompt_list = [] qid_list = [] for item in data_list: question = item['question_content'].strip() if item['starter_code'] != "": question += "\n\n" + "Solve the problem starting with the provided function header.\n\nFunction header:\n" + "```\n" + \ item['starter_code'] + "\n```" question += "\n\n" + code_instruction_hasstartercode else: question += "\n\n" + code_instruction_nostartercode final_prompt = question prompt_list.append(final_prompt) qid_list.append(item['question_id']) return prompt_list, qid_list def preprocess_mbpp_chatml_template(data_file): """Preprocess MBPP (Mostly Basic Python Problems) dataset with ChatML template. Args: data_file: Path to MBPP JSON file Returns: tuple: (prompt_list, qid_list) - prompt_list: Formatted code generation prompts with ChatML template - qid_list: Task IDs """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" with open(data_file, "r") as f: data_dict = json.load(f) prompt_list = [] qid_list = [] for key, value in data_dict.items(): qid_list.append(key) question = value.get('text', value.get('prompt', '')) final_prompt = instruction + "<|im_start|>user\n" + question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list, qid_list def preprocess_mmlu_stem_chatml_template(data_file): """Preprocess MMLU STEM subset with ChatML template. Args: data_file: Path to MMLU STEM JSON file Returns: list: Formatted prompts with ChatML template """ instruction = "<|im_start|>system\nYou are a helpful and harmless assistant. You should think step-by-step.<|im_end|>\n" with open(data_file, "r") as f: data_list = json.load(f) prompt_list = [] for item in data_list: question = item['question'].strip() choices = item.get('choices', []) final_question = question + "\n" if choices: for i, choice in enumerate(choices): final_question += f"({chr(65+i)}) {choice}\n" final_prompt = instruction + "<|im_start|>user\n" + final_question + "<|im_end|>\n<|im_start|>assistant\n\n" prompt_list.append(final_prompt) return prompt_list