bench_other.py 4.89 KB
Newer Older
Lianmin Zheng's avatar
Lianmin Zheng committed
1
2
3
4
5
import argparse
import asyncio
import json
import os
import time
Liangsheng Yin's avatar
Liangsheng Yin committed
6
from concurrent.futures import ThreadPoolExecutor
Lianmin Zheng's avatar
Lianmin Zheng committed
7
8
9
10
11
12

import numpy as np
import pandas as pd
import tiktoken
from tqdm import tqdm

Liangsheng Yin's avatar
Liangsheng Yin committed
13
from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate
Lianmin Zheng's avatar
Lianmin Zheng committed
14
15
16
17
18
19
20
21
22
23
24
25
26

choices = ["A", "B", "C", "D"]

tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")


def format_subject(subject):
    l = subject.split("_")
    s = ""
    for entry in l:
        s += " " + entry
    return s

Liangsheng Yin's avatar
Liangsheng Yin committed
27

Lianmin Zheng's avatar
Lianmin Zheng committed
28
29
30
31
def format_example(df, idx, include_answer=True):
    prompt = df.iloc[idx, 0]
    k = df.shape[1] - 2
    for j in range(k):
Liangsheng Yin's avatar
Liangsheng Yin committed
32
        prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1])
Lianmin Zheng's avatar
Lianmin Zheng committed
33
34
35
36
37
    prompt += "\nAnswer:"
    if include_answer:
        prompt += " {}\n\n".format(df.iloc[idx, k + 1])
    return prompt

Liangsheng Yin's avatar
Liangsheng Yin committed
38

Lianmin Zheng's avatar
Lianmin Zheng committed
39
def gen_prompt(train_df, subject, k=-1):
Liangsheng Yin's avatar
Liangsheng Yin committed
40
41
42
    prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format(
        format_subject(subject)
    )
Lianmin Zheng's avatar
Lianmin Zheng committed
43
44
45
46
47
48
49
    if k == -1:
        k = train_df.shape[0]
    for i in range(k):
        prompt += format_example(train_df, i)
    return prompt


Liangsheng Yin's avatar
Liangsheng Yin committed
50
def evaluate(args, subject, dev_df, test_df, call_generate):
Lianmin Zheng's avatar
Lianmin Zheng committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
    prompts = []
    labels = []

    # Construct prompts
    k = args.ntrain
    train_prompt = gen_prompt(dev_df, subject, k)
    while len(tokenizer.encode(train_prompt)) > 1536:
        k -= 1
        train_prompt = gen_prompt(dev_df, subject, k)

    for i in range(test_df.shape[0]):
        prompt_end = format_example(test_df, i, include_answer=False)
        prompt = train_prompt + prompt_end
        prompts.append(prompt)

Liangsheng Yin's avatar
Liangsheng Yin committed
66
        label = test_df.iloc[i, test_df.shape[1] - 1]
Lianmin Zheng's avatar
Lianmin Zheng committed
67
68
69
70
71
72
73
74
75
        labels.append(label)

    preds = [None] * len(prompts)
    max_tokens = 1

    # Run requests
    if args.backend != "lmql":
        # Use thread pool
        def get_one_answer(i):
Liangsheng Yin's avatar
Liangsheng Yin committed
76
            pred = call_generate(prompts[i], temperature=0, max_tokens=max_tokens)
Lianmin Zheng's avatar
Lianmin Zheng committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
            preds[i] = pred.strip()[0]

        tic = time.time()
        if args.parallel == 1:
            for i in range(len(prompts)):
                get_one_answer(i)
        else:
            with ThreadPoolExecutor(args.parallel) as executor:
                executor.map(get_one_answer, list(range(len(prompts))))
    else:
        # Use asyncio
        async def batched_call(batch_size):
            for i in range(0, len(prompts), batch_size):
                tasks = []
Liangsheng Yin's avatar
Liangsheng Yin committed
91
92
                for p in prompts[i : i + batch_size]:
                    tasks.append(call_generate(p, temperature=0, max_tokens=max_tokens))
Lianmin Zheng's avatar
Lianmin Zheng committed
93
94
                rets = await asyncio.gather(*tasks)
                for j in range(len(rets)):
Liangsheng Yin's avatar
Liangsheng Yin committed
95
                    preds[i + j] = rets[j].strip()[0]
Lianmin Zheng's avatar
Lianmin Zheng committed
96
97
98
99
100
101
102
103
104
105

        tic = time.time()
        asyncio.run(batched_call(batch_size=args.parallel))
    latency = time.time() - tic

    # Compute accuracy
    cors = [pred == label for pred, label in zip(preds, labels)]
    acc = np.mean(cors)
    cors = np.array(cors)

Liangsheng Yin's avatar
Liangsheng Yin committed
106
107
108
109
110
    print(
        "Average accuracy {:.3f}, latency {:.2f}, #q: {} - {}".format(
            acc, latency, len(prompts), subject
        )
    )
Lianmin Zheng's avatar
Lianmin Zheng committed
111
112
113
114
115

    return cors, acc, latency


def main(args):
Liangsheng Yin's avatar
Liangsheng Yin committed
116
117
118
119
120
121
122
    subjects = sorted(
        [
            f.split("_test.csv")[0]
            for f in os.listdir(os.path.join(args.data_dir, "test"))
            if "_test.csv" in f
        ]
    )
Lianmin Zheng's avatar
Lianmin Zheng committed
123
124
125
126
127

    all_cors = []
    all_latencies = []
    num_requests = 0

Liangsheng Yin's avatar
Liangsheng Yin committed
128
129
130
    # Select backend
    call_generate = get_call_generate(args)

Liangsheng Yin's avatar
Liangsheng Yin committed
131
132
133
134
135
136
137
    for subject in tqdm(subjects[: args.nsub]):
        dev_df = pd.read_csv(
            os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None
        )[: args.ntrain]
        test_df = pd.read_csv(
            os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
138

Liangsheng Yin's avatar
Liangsheng Yin committed
139
        cors, acc, latency = evaluate(args, subject, dev_df, test_df, call_generate)
Lianmin Zheng's avatar
Lianmin Zheng committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
        all_cors.append(cors)
        all_latencies.append(latency)
        num_requests += len(test_df)

    total_latency = np.sum(all_latencies)
    print("Total latency: {:.3f}".format(total_latency))

    weighted_acc = np.mean(np.concatenate(all_cors))
    print("Average accuracy: {:.3f}".format(weighted_acc))

    # Write results
    with open(args.result_file, "a") as fout:
        value = {
            "task": "mmlu",
            "backend": args.backend,
            "num_gpus": 1,
            "latency": round(total_latency, 3),
            "accuracy": round(weighted_acc, 3),
            "num_requests": num_requests,
            "other": {
                "nsub": args.nsub,
                "parallel": args.parallel,
Liangsheng Yin's avatar
Liangsheng Yin committed
162
            },
Lianmin Zheng's avatar
Lianmin Zheng committed
163
164
165
166
167
168
169
170
171
172
173
        }
        fout.write(json.dumps(value) + "\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ntrain", type=int, default=5)
    parser.add_argument("--data_dir", type=str, default="data")
    parser.add_argument("--nsub", type=int, default=60)
    args = add_common_other_args_and_parse(parser)
    main(args)