File size: 2,627 Bytes
31dc8dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import json
import os
from datasets import load_dataset
def format_gsm8k_example_to_messages(example):
"""
Converts a single example from the gsm8k dataset into the messages format.
Args:
example (dict): A single sample from the dataset, containing 'question' and 'answer' keys.
Returns:
dict: A dictionary containing a 'messages' key, whose value is a list
of dictionaries representing the user and assistant dialogue.
"""
# 'question' corresponds to the 'user' role
# 'answer' corresponds to the 'assistant' role
return {
"messages": [
{"role": "user", "content": example["question"]},
{"role": "assistant", "content": example["answer"]}
]
}
def main():
# 1. Load the gsm8k dataset (main config) from the Hugging Face Hub
print("Loading gsm8k (main) dataset...")
# The 'main' configuration contains clean question-answer pairs
raw_dataset = load_dataset("openai/gsm8k", "main")
print("Dataset loaded successfully.")
print(f"Original dataset structure: {raw_dataset}")
print(f"First example from the original training set: {raw_dataset['train'][0]}")
# 2. Use the .map() method to convert the dataset to the messages format
# remove_columns will drop the old 'question' and 'answer' columns, keeping only the new 'messages' column
print("\nConverting dataset to messages format...")
transformed_dataset = raw_dataset.map(
format_gsm8k_example_to_messages,
remove_columns=["question", "answer"]
)
print("Conversion complete.")
print(f"Transformed dataset structure: {transformed_dataset}")
print(f"First example from the transformed training set: {transformed_dataset['train'][0]}")
# 3. Save the processed dataset to JSONL files
print("\nSaving to JSONL files...")
# Ensure the output directory exists
output_dir = "./gsm8k_datasets"
os.makedirs(output_dir, exist_ok=True)
# Iterate through each split of the dataset (e.g., 'train' and 'test')
for split, dataset_split in transformed_dataset.items():
output_filename = os.path.join(output_dir, f"gsm8k_{split}.jsonl")
with open(output_filename, 'w', encoding='utf-8') as f:
for record in dataset_split:
# Convert each dictionary to a JSON string and write it to the file, followed by a newline
json_string = json.dumps(record, ensure_ascii=False)
f.write(json_string + '\n')
print(f"Saved the '{split}' split to: {output_filename}")
if __name__ == "__main__":
main()
|