-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathrun.py
122 lines (102 loc) · 4.33 KB
/
run.py
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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
from ada_leval.smp import *
from ada_leval.util import *
from ada_leval.api import OpenAIWrapper
from ada_leval.dataset import StackSelect, TextSort
RESULT_FILE = 'result.json'
if not osp.exists(RESULT_FILE):
dump({}, RESULT_FILE)
settings = ['1k', '2k', '4k', '8k', '16k', '32k', '64k', '128k']
datasets = [f'stackselect_{k}' for k in settings + ['6k', '12k']] + [f'textsort_{k}' for k in settings]
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, nargs='+', required=True, choices=datasets)
parser.add_argument('--model', type=str, required=True, choices=['internlm2-7b', 'internlm2-20b', 'gpt-4-0125'])
parser.add_argument('--mode', type=str, default='all', choices=['infer', 'all'])
parser.add_argument('--nproc', type=int, default=4)
args = parser.parse_args()
return args
def build_model(m, setting=None):
if 'internlm2' in m:
session_len = 160000
from lmdeploy import pipeline, TurbomindEngineConfig
backend_config = TurbomindEngineConfig(rope_scaling_factor=2.0, session_len=session_len)
if m == 'gpt-4-0125':
model = OpenAIWrapper('gpt-4-0125-preview')
elif m == 'internlm2-7b':
model = pipeline('internlm/internlm2-chat-7b', backend_config=backend_config)
elif m == 'internlm2-20b':
model = pipeline('internlm/internlm2-chat-20b', backend_config=backend_config)
return model
import tiktoken
ENC = tiktoken.encoding_for_model('gpt-4')
def get_token_length(prompt):
return len(ENC.encode(prompt))
def main():
rank, world_size = get_rank_and_world_size()
if world_size > 1:
import torch
import torch.distributed as dist
torch.cuda.set_device(rank)
dist.init_process_group(backend='nccl')
if rank == 0:
os.makedirs('results', exist_ok=True)
args = parse_args()
model_name = args.model
model = build_model(args.model)
for dname in args.data:
d, setting = dname.split('_')
dataset_mode = 'less' if getattr(model, 'is_api', False) else 'normal'
if d == 'stackselect':
dataset = StackSelect(setting=setting, mode=dataset_mode)
elif d == 'textsort':
dataset = TextSort(setting=setting, mode=dataset_mode)
lt = len(dataset)
prompts = [dataset.build_prompt(i) for i in range(lt)]
meta = dataset.get_meta()
indices = list(meta['index'])
out_file = f'results/{model_name}_{dname}.pkl'
res = {} if not osp.exists(out_file) else load(out_file)
tups = [(i, p) for i, p in zip(indices, prompts) if i not in res]
if len(tups):
if getattr(model, 'is_api', False):
res = track_progress_rich(
model.generate,
[x[1] for x in tups],
nproc=args.nproc,
chunksize=args.nproc,
save=out_file,
keys=[x[0] for x in tups])
else:
sub_tups = tups[rank::world_size]
sub_out_file = f'results/{model_name}_{dname}_{rank}.pkl'
sub_res = {}
import torch
with torch.no_grad():
for t in tqdm(sub_tups):
index, prompt = t
sub_res[index] = model(prompt).text
dump(sub_res, sub_out_file)
if world_size > 1:
dist.barrier()
if rank == 0:
res = {}
for i in range(world_size):
sub_out_file = f'results/{model_name}_{dname}_{i}.pkl'
if osp.exists(sub_out_file):
res.update(load(sub_out_file))
if osp.exists(out_file):
res.update(load(out_file))
dump(res, out_file)
res = load(out_file)
meta['prediction'] = [res[k] for k in meta['index']]
dump(meta, f'results/{model_name}_{dname}.xlsx')
if args.mode == 'all':
results = load(RESULT_FILE)
acc = dataset.evaluate(meta)
results[f'{model_name}_{dname}'] = acc
dump(results, RESULT_FILE)
if world_size > 1:
dist.barrier()
os.system(f"rm {f'results/{model_name}_{dname}_{rank}.pkl'}")
if __name__ == '__main__':
main()