File size: 3,347 Bytes
2a26d3b |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
import json
import random
import datasets
from datasets import load_dataset, Dataset, concatenate_datasets
from utils import load_jsonl, lower_keys
def load_data(data_name, split, data_dir="./data"):
data_file = f"{data_dir}/{data_name}/{split}.jsonl"
if os.path.exists(data_file):
examples = list(load_jsonl(data_file))
else:
if data_name == "math":
dataset = load_dataset(
"competition_math",
split=split,
name="main",
cache_dir=f"{data_dir}/temp",
)
elif data_name == "gsm8k":
dataset = load_dataset(data_name, split=split)
elif data_name == "svamp":
# evaluate on training set + test set
dataset = load_dataset("ChilleD/SVAMP", split="train")
dataset = concatenate_datasets(
[dataset, load_dataset("ChilleD/SVAMP", split="test")]
)
elif data_name == "asdiv":
dataset = load_dataset("EleutherAI/asdiv", split="validation")
dataset = dataset.filter(
lambda x: ";" not in x["answer"]
) # remove multi-answer examples
elif data_name == "mawps":
examples = []
# four sub-tasks
for data_name in ["singleeq", "singleop", "addsub", "multiarith"]:
sub_examples = list(load_jsonl(f"{data_dir}/mawps/{data_name}.jsonl"))
for example in sub_examples:
example["type"] = data_name
examples.extend(sub_examples)
dataset = Dataset.from_list(examples)
elif data_name == "mmlu_stem":
dataset = load_dataset("hails/mmlu_no_train", "all", split="test")
# only keep stem subjects
stem_subjects = [
"abstract_algebra",
"astronomy",
"college_biology",
"college_chemistry",
"college_computer_science",
"college_mathematics",
"college_physics",
"computer_security",
"conceptual_physics",
"electrical_engineering",
"elementary_mathematics",
"high_school_biology",
"high_school_chemistry",
"high_school_computer_science",
"high_school_mathematics",
"high_school_physics",
"high_school_statistics",
"machine_learning",
]
dataset = dataset.rename_column("subject", "type")
dataset = dataset.filter(lambda x: x["type"] in stem_subjects)
elif data_name == "carp_en":
dataset = load_jsonl(f"{data_dir}/carp_en/test.jsonl")
else:
raise NotImplementedError(data_name)
examples = list(dataset)
examples = [lower_keys(example) for example in examples]
dataset = Dataset.from_list(examples)
os.makedirs(f"{data_dir}/{data_name}", exist_ok=True)
dataset.to_json(data_file)
# add 'idx' in the first column
if "idx" not in examples[0]:
examples = [{"idx": i, **example} for i, example in enumerate(examples)]
# dedepulicate & sort
examples = sorted(examples, key=lambda x: x["idx"])
return examples
|