# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import argparse import csv import io import json import os import random import tarfile import urllib.request from pathlib import Path from typing import Any, Dict, List, Tuple from datasets import load_dataset random.seed(42) subcategories = { "abstract_algebra": ["math"], "anatomy": ["health"], "astronomy": ["physics"], "business_ethics": ["business"], "clinical_knowledge": ["health"], "college_biology": ["biology"], "college_chemistry": ["chemistry"], "college_computer_science": ["computer science"], "college_mathematics": ["math"], "college_medicine": ["health"], "college_physics": ["physics"], "computer_security": ["computer science"], "conceptual_physics": ["physics"], "econometrics": ["economics"], "electrical_engineering": ["engineering"], "elementary_mathematics": ["math"], "formal_logic": ["philosophy"], "global_facts": ["other"], "high_school_biology": ["biology"], "high_school_chemistry": ["chemistry"], "high_school_computer_science": ["computer science"], "high_school_european_history": ["history"], "high_school_geography": ["geography"], "high_school_government_and_politics": ["politics"], "high_school_macroeconomics": ["economics"], "high_school_mathematics": ["math"], "high_school_microeconomics": ["economics"], "high_school_physics": ["physics"], "high_school_psychology": ["psychology"], "high_school_statistics": ["math"], "high_school_us_history": ["history"], "high_school_world_history": ["history"], "human_aging": ["health"], "human_sexuality": ["culture"], "international_law": ["law"], "jurisprudence": ["law"], "logical_fallacies": ["philosophy"], "machine_learning": ["computer science"], "management": ["business"], "marketing": ["business"], "medical_genetics": ["health"], "miscellaneous": ["other"], "moral_disputes": ["philosophy"], "moral_scenarios": ["philosophy"], "nutrition": ["health"], "philosophy": ["philosophy"], "prehistory": ["history"], "professional_accounting": ["other"], "professional_law": ["law"], "professional_medicine": ["health"], "professional_psychology": ["psychology"], "public_relations": ["politics"], "security_studies": ["politics"], "sociology": ["culture"], "us_foreign_policy": ["politics"], "virology": ["health"], "world_religions": ["philosophy"], } def process_mmlu_dataset(split) -> List[Dict[str, Any]]: """Process the MMLU dataset and return a list of formatted entries.""" data_url = "https://people.eecs.berkeley.edu/~hendrycks/data.tar" data_dir = Path(__file__).absolute().parent data_file = str(data_dir / f"data.tar") data_dir.mkdir(exist_ok=True) output_file = str(data_dir / f"{split}.jsonl") urllib.request.urlretrieve(data_url, data_file) result = {} column_names = ["question", "A", "B", "C", "D", "expected_answer"] with tarfile.open(data_file, 'r') as tar: # List all members of the tar file members = tar.getmembers() # Filter for CSV files in the 'data/test' directory csv_files = [ member for member in members if member.name.startswith(f'data/{split}/') and member.name.endswith('.csv') ] for csv_file in csv_files: # Extract the file name without the path file_name = os.path.basename(csv_file.name) # Read the CSV file content file_content = tar.extractfile(csv_file) if file_content is not None: # Decode bytes to string content_str = io.TextIOWrapper(file_content, encoding='utf-8') # Use csv to read the CSV content without a header csv_reader = csv.reader(content_str) # Convert CSV data to list of dictionaries with specified column names csv_data = [] for row in csv_reader: if len(row) == len(column_names): csv_data.append(dict(zip(column_names, row))) else: print(f"Warning: Skipping row in {file_name} due to incorrect number of columns") # Add to result dictionary result[file_name.rsplit('_', 1)[0]] = csv_data # Define the allowed supercategories chosen_subcategories = {"math", "health", "physics", "biology", "chemistry", "computer science", "engineering"} # Filter the result dictionary filtered_result = { k: v for k, v in result.items() if any(supercat in chosen_subcategories for supercat in subcategories.get(k, [])) } output_data = [] skipped = 0 for category, question_list in filtered_result.items(): for item in question_list: question = item['question'] choice_1 = item['A'] choice_2 = item['B'] choice_3 = item['C'] choice_4 = item['D'] answer = item['expected_answer'] if not question or not answer or not choice_1 or not choice_2 or not choice_3 or not choice_4: skipped += 1 continue entry = { 'Question': question, 'Choice 1': choice_1, 'Choice 2': choice_2, 'Choice 3': choice_3, 'Choice 4': choice_4, 'Answer': answer, 'Subject': category, } output_data.append(entry) return output_data, skipped def process_gpqa_dataset() -> List[Dict[str, Any]]: """Process the GPQA dataset and return a list of formatted entries.""" dataset = load_dataset("Idavidrein/gpqa", "gpqa_main") output_data = [] skipped = 0 for split in dataset.keys(): for item in dataset[split]: question = item['Question'] correct_answer = item['Correct Answer'] incorrect_1 = item['Incorrect Answer 1'] incorrect_2 = item['Incorrect Answer 2'] incorrect_3 = item['Incorrect Answer 3'] if not question or not correct_answer or not incorrect_1 or not incorrect_2 or not incorrect_3: skipped += 1 continue answers = [(correct_answer, True), (incorrect_1, False), (incorrect_2, False), (incorrect_3, False)] random.shuffle(answers) choices = {} correct_letter = None for i, (answer_text, is_correct) in enumerate(answers): letter = chr(65 + i) choices[f'Choice {i+1}'] = answer_text if is_correct: correct_letter = letter entry = { 'Question': question, 'Choice 1': choices['Choice 1'], 'Choice 2': choices['Choice 2'], 'Choice 3': choices['Choice 3'], 'Choice 4': choices['Choice 4'], 'Answer': correct_letter, } output_data.append(entry) return output_data, skipped def process_gpqa_diamond_dataset() -> List[Dict[str, Any]]: """Process the GPQA Diamond dataset and return a list of formatted entries.""" dataset = load_dataset("Idavidrein/gpqa", "gpqa_diamond") output_data = [] skipped = 0 for split in dataset.keys(): for item in dataset[split]: question = item['Question'] correct_answer = item['Correct Answer'] incorrect_1 = item['Incorrect Answer 1'] incorrect_2 = item['Incorrect Answer 2'] incorrect_3 = item['Incorrect Answer 3'] if not question or not correct_answer or not incorrect_1 or not incorrect_2 or not incorrect_3: skipped += 1 continue answers = [(correct_answer, True), (incorrect_1, False), (incorrect_2, False), (incorrect_3, False)] random.shuffle(answers) choices = {} correct_letter = None for i, (answer_text, is_correct) in enumerate(answers): letter = chr(65 + i) choices[f'Choice {i+1}'] = answer_text if is_correct: correct_letter = letter entry = { 'Question': question, 'Choice 1': choices['Choice 1'], 'Choice 2': choices['Choice 2'], 'Choice 3': choices['Choice 3'], 'Choice 4': choices['Choice 4'], 'Answer': correct_letter, } output_data.append(entry) return output_data, skipped def write_to_jsonl(data: List[Dict[str, Any]], filename: str) -> None: """Write the processed data to a JSONL file.""" with open(filename, 'w') as f: for entry in data: f.write(json.dumps(entry) + '\n') def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser(description='Convert datasets to JSONL format') parser.add_argument( '--datasets', nargs='+', choices=['mmlu', 'gpqa', 'gpqa_diamond', 'all'], default=['all'], help='Datasets to process (default: all)', ) return parser.parse_args() def main(): args = parse_args() # If 'all' is selected, process all datasets if 'all' in args.datasets: args.datasets = ['mmlu', 'gpqa', 'gpqa_diamond'] # Process selected datasets if 'mmlu' in args.datasets: for split in ["test", "val"]: mmlu_data, mmlu_skipped = process_mmlu_dataset(split) write_to_jsonl(mmlu_data, f'mmlu_dataset_{split}.jsonl') print(f"Converted {len(mmlu_data)} {split} questions to mmlu_dataset_{split}.jsonl") if mmlu_skipped > 0: print(f"Skipped {mmlu_skipped} MMLU entries due to missing data") print("\nMMLU Sample entry:") print(json.dumps(mmlu_data[0], indent=2)) if 'gpqa' in args.datasets: gpqa_data, gpqa_skipped = process_gpqa_dataset() write_to_jsonl(gpqa_data, 'gpqa_dataset.jsonl') print(f"\nConverted {len(gpqa_data)} questions to gpqa_dataset.jsonl") if gpqa_skipped > 0: print(f"Skipped {gpqa_skipped} GPQA entries due to missing data") print("\nGPQA Sample entry:") print(json.dumps(gpqa_data[0], indent=2)) if 'gpqa_diamond' in args.datasets: gpqa_diamond_data, gpqa_diamond_skipped = process_gpqa_diamond_dataset() write_to_jsonl(gpqa_diamond_data, 'gpqa_diamond_dataset.jsonl') print(f"\nConverted {len(gpqa_diamond_data)} questions to gpqa_diamond_dataset.jsonl") if gpqa_diamond_skipped > 0: print(f"Skipped {gpqa_diamond_skipped} GPQA Diamond entries due to missing data") print("\nGPQA Diamond Sample entry:") print(json.dumps(gpqa_diamond_data[0], indent=2)) if __name__ == "__main__": main()