Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
gen_all_splits.py 2.67 KB
Newer Older
haileyschoelkopf's avatar
add yaml-genning script + WIP base yaml  
haileyschoelkopf committed
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105 
"""
Take in a YAML, and output all other splits with this YAML
"""
import os
import yaml
import argparse

from tqdm import tqdm

from lm_eval import utils
from lm_eval.logger import eval_logger

SUBJECTS = [
    # "abstract_algebra",
    "anatomy",
    "astronomy",
    "business_ethics",
    "clinical_knowledge",
    "college_biology",
    "college_chemistry",
    "college_computer_science",
    "college_mathematics",
    "college_medicine",
    "college_physics",
    "computer_security",
    "conceptual_physics",
    "econometrics",
    "electrical_engineering",
    "elementary_mathematics",
    "formal_logic",
    "global_facts",
    "high_school_biology",
    "high_school_chemistry",
    "high_school_computer_science",
    "high_school_european_history",
    "high_school_geography",
    "high_school_government_and_politics",
    "high_school_macroeconomics",
    "high_school_mathematics",
    "high_school_microeconomics",
    "high_school_physics",
    "high_school_psychology",
    "high_school_statistics",
    "high_school_us_history",
    "high_school_world_history",
    "human_aging",
    "human_sexuality",
    "international_law",
    "jurisprudence",
    "logical_fallacies",
    "machine_learning",
    "management",
    "marketing",
    "medical_genetics",
    "miscellaneous",
    "moral_disputes",
    "moral_scenarios",
    "nutrition",
    "philosophy",
    "prehistory",
    "professional_accounting",
    "professional_law",
    "professional_medicine",
    "professional_psychology",
    "public_relations",
    "security_studies",
    "sociology",
    "us_foreign_policy",
    "virology",
    "world_religions",
]


def parse_args():
    parser = argparse.ArgumentParser()
    # parser.add_argument("--benchmark_name", required=True)
    parser.add_argument("--base_yaml_path", required=True)
    parser.add_argument(
        "--task_save_path", default="lm_eval/tasks/mmlu/hendrycks_test_original"
    )
    return parser.parse_args()


if __name__ == "__main__":

    args = parse_args()

    # get filename of base_yaml so we can `"include": ` it in our other YAMLs.
    base_yaml_name = os.path.split(args.base_yaml_path)[-1]
    with open(args.base_yaml_path) as f:
        base_yaml = yaml.full_load(f)
    print(base_yaml)

    for subject in tqdm(SUBJECTS):

        yaml_dict = {
            "include": base_yaml_name,
            "task": base_yaml["task"].strip("abstract_algebra") + "subject",
            "dataset_name": subject,
        }

        file_save_path = args.task_save_path + f"_{subject}.yaml"
        eval_logger.info(f"Saving yaml for subset {subject} to {file_save_path}")
        with open(file_save_path, "w") as yaml_file:
            yaml.dump(yaml_dict, yaml_file)