main.py 4.41 KB
Newer Older
lintangsutawika's avatar
lintangsutawika committed
1
import os
lintangsutawika's avatar
lintangsutawika committed
2
import re
Jason Phang's avatar
Jason Phang committed
3
import json
4
import fnmatch
lintangsutawika's avatar
lintangsutawika committed
5
import jsonlines
lintangsutawika's avatar
lintangsutawika committed
6
import argparse
FarzanehNakhaee's avatar
FarzanehNakhaee committed
7
import logging
Leo Gao's avatar
Leo Gao committed
8

9
from lm_eval import evaluator, utils
10
from lm_eval.api.registry import ALL_TASKS
lintangsutawika's avatar
lintangsutawika committed
11
from lm_eval.logger import eval_logger
Jason Phang's avatar
lib  
Jason Phang committed
12

13
os.environ["TOKENIZERS_PARALLELISM"] = "false"
14

Fabrizio Milo's avatar
Fabrizio Milo committed
15

Jason Phang's avatar
Jason Phang committed
16
17
def parse_args():
    parser = argparse.ArgumentParser()
Fabrizio Milo's avatar
Fabrizio Milo committed
18
19
    parser.add_argument("--model", required=True)
    parser.add_argument("--model_args", default="")
lintangsutawika's avatar
lintangsutawika committed
20
21
22
    parser.add_argument(
        "--tasks", default=None, choices=utils.MultiChoice(sorted(ALL_TASKS))
    )
23
    parser.add_argument("--config", default=None)
Fabrizio Milo's avatar
Fabrizio Milo committed
24
    parser.add_argument("--num_fewshot", type=int, default=0)
25
    parser.add_argument("--batch_size", type=int, default=1)
lintangsutawika's avatar
lintangsutawika committed
26
27
28
29
30
31
    parser.add_argument(
        "--max_batch_size",
        type=int,
        default=None,
        help="Maximal batch size to try with --batch_size auto",
    )
Fabrizio Milo's avatar
Fabrizio Milo committed
32
33
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--output_path", default=None)
lintangsutawika's avatar
lintangsutawika committed
34
35
36
37
38
39
40
    parser.add_argument(
        "--limit",
        type=float,
        default=None,
        help="Limit the number of examples per task. "
        "If <1, limit is a percentage of the total number of examples.",
    )
41
    parser.add_argument("--data_sampling", type=float, default=None)
Fabrizio Milo's avatar
Fabrizio Milo committed
42
43
44
    parser.add_argument("--no_cache", action="store_true")
    parser.add_argument("--decontamination_ngrams_path", default=None)
    parser.add_argument("--check_integrity", action="store_true")
45
46
    parser.add_argument("--write_out", action="store_true", default=False)
    parser.add_argument("--output_base_path", type=str, default=None)
Jason Phang's avatar
Jason Phang committed
47
48
    return parser.parse_args()

Fabrizio Milo's avatar
Fabrizio Milo committed
49

50
def main():
Jason Phang's avatar
Jason Phang committed
51
    args = parse_args()
Fabrizio Milo's avatar
Fabrizio Milo committed
52

Leo Gao's avatar
Leo Gao committed
53
    if args.limit:
lintangsutawika's avatar
lintangsutawika committed
54
55
56
        eval_logger.warning(
            " --limit SHOULD ONLY BE USED FOR TESTING."
            "REAL METRICS SHOULD NOT BE COMPUTED USING LIMIT."
Fabrizio Milo's avatar
Fabrizio Milo committed
57
        )
Leo Gao's avatar
Leo Gao committed
58

59
    if args.tasks is None:
60
        task_names = ALL_TASKS
Jason Phang's avatar
Jason Phang committed
61
    else:
62
63
        if os.path.isdir(args.tasks):
            import glob
64
65

            task_names = []
66
67
            yaml_path = os.path.join(args.tasks, "*.yaml")
            for yaml_file in glob.glob(yaml_path):
lintangsutawika's avatar
lintangsutawika committed
68
                config = utils.load_yaml_config(yaml_file)
69
70
                task_names.append(config)
        else:
71
            tasks_list = args.tasks.split(",")
72
            task_names = utils.pattern_match(tasks_list, ALL_TASKS)
73
74
            for task in [task for task in tasks_list if task not in task_names]:
                if os.path.isfile(task):
lintangsutawika's avatar
lintangsutawika committed
75
                    config = utils.load_yaml_config(task)
76
                    task_names.append(config)
lintangsutawika's avatar
lintangsutawika committed
77

lintangsutawika's avatar
lintangsutawika committed
78
    eval_logger.info(f"Selected Tasks: {task_names}")
79

80
81
82
83
84
85
    results = evaluator.simple_evaluate(
        model=args.model,
        model_args=args.model_args,
        tasks=task_names,
        num_fewshot=args.num_fewshot,
        batch_size=args.batch_size,
86
        max_batch_size=args.max_batch_size,
87
        device=args.device,
88
        no_cache=args.no_cache,
89
90
91
        limit=args.limit,
        decontamination_ngrams_path=args.decontamination_ngrams_path,
        check_integrity=args.check_integrity,
92
93
        write_out=args.write_out,
        output_base_path=args.output_base_path,
94
    )
95

96
    if results is not None:
lintangsutawika's avatar
lintangsutawika committed
97
98
        samples = results.pop("samples")

99
100
101
        dumped = json.dumps(results, indent=2)
        print(dumped)

102
103
        batch_sizes = ",".join(map(str, results["config"]["batch_sizes"]))

104
        if args.output_path:
105
            os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
lintangsutawika's avatar
lintangsutawika committed
106

107
108
109
            with open(args.output_path, "w") as f:
                f.write(dumped)

lintangsutawika's avatar
lintangsutawika committed
110
111
112
113
114
115
116
117
118
119
120
121
122
123
            for task_name, config in results["configs"].items():
                output_name = "{}_{}".format(
                    re.sub("/", "__", args.model_args), task_name
                )
                if os.path.isdir(args.output_path):
                    filename = f"./{args.output_path}/{output_name}.jsonl"
                elif os.path.isfile(args.output_path):
                    filename = (
                        f"./{os.path.dirname(args.output_path)}/{output_name}.jsonl"
                    )

                with jsonlines.open(filename, "w") as f:
                    f.write_all(samples[task_name])

124
        print(
125
126
            f"{args.model} ({args.model_args}), limit: {args.limit}, num_fewshot: {args.num_fewshot}, "
            f"batch_size: {args.batch_size}{f' ({batch_sizes})' if batch_sizes else ''}"
127
128
        )
        print(evaluator.make_table(results))
Jason Phang's avatar
lib  
Jason Phang committed
129

130

Jason Phang's avatar
Jason Phang committed
131
if __name__ == "__main__":
Jason Phang's avatar
lib  
Jason Phang committed
132
    main()