|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
|
Preprocess the MATH-lighteval dataset to parquet format |
|
|
""" |
|
|
|
|
|
import argparse |
|
|
import os |
|
|
|
|
|
import datasets |
|
|
|
|
|
from verl.utils.hdfs_io import copy, makedirs |
|
|
from verl.utils.reward_score.math import last_boxed_only_string, remove_boxed |
|
|
|
|
|
|
|
|
def extract_solution(solution_str): |
|
|
return remove_boxed(last_boxed_only_string(solution_str)) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--local_dir", default="~/data/math") |
|
|
parser.add_argument("--hdfs_dir", default=None) |
|
|
|
|
|
args = parser.parse_args() |
|
|
|
|
|
|
|
|
|
|
|
data_source = "DigitalLearningGmbH/MATH-lighteval" |
|
|
print(f"Loading the {data_source} dataset from huggingface...", flush=True) |
|
|
dataset = datasets.load_dataset(data_source, trust_remote_code=True) |
|
|
|
|
|
train_dataset = dataset["train"] |
|
|
test_dataset = dataset["test"] |
|
|
|
|
|
instruction_following = "Let's think step by step and output the final answer within \\boxed{}." |
|
|
|
|
|
|
|
|
def make_map_fn(split): |
|
|
def process_fn(example, idx): |
|
|
question = example.pop("problem") |
|
|
|
|
|
question = question + " " + instruction_following |
|
|
|
|
|
answer = example.pop("solution") |
|
|
solution = extract_solution(answer) |
|
|
data = { |
|
|
"data_source": data_source, |
|
|
"prompt": [{"role": "user", "content": question}], |
|
|
"ability": "math", |
|
|
"reward_model": {"style": "rule", "ground_truth": solution}, |
|
|
"extra_info": {"split": split, "index": idx}, |
|
|
} |
|
|
return data |
|
|
|
|
|
return process_fn |
|
|
|
|
|
train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) |
|
|
test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) |
|
|
|
|
|
local_dir = args.local_dir |
|
|
hdfs_dir = args.hdfs_dir |
|
|
|
|
|
train_dataset.to_parquet(os.path.join(local_dir, "train.parquet")) |
|
|
test_dataset.to_parquet(os.path.join(local_dir, "test.parquet")) |
|
|
|
|
|
if hdfs_dir is not None: |
|
|
makedirs(hdfs_dir) |
|
|
|
|
|
copy(src=local_dir, dst=hdfs_dir) |
|
|
|