CodeEditerBench/batch_api_inference.py at main · jiamin0630/CodeEditerBench · GitHub

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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
import json
import os
import time
from tqdm import tqdm
from api_handler import api_handler
from api_inference import (generate_prompt_debug, generate_prompt_translate,
                        generate_prompt_polishment, generate_prompt_switch,
                        generate_cot_prompt_debug, generate_cot_prompt_translate,
                        generate_cot_prompt_polishment, generate_cot_prompt_switch,
                        read_jsonl_file)
import argparse
import jsonlines

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default="gpt-35-turbo", help="Model to use: gpt-35-turbo or gpt-4")
    parser.add_argument("--task", type=str, default="debug", help="Task: debug, translate, polishment, or switch")
    parser.add_argument("--prompt_type", type=str, default="zero", help="Prompt type: zero, three, or cot")
    parser.add_argument("--dataset_type", type=str, default="primary", choices=["primary", "plus"],
                        help="Dataset type: primary or plus (only affects input data)")
    parser.add_argument("--start_idx", type=int, default=0, help="Start index")
    parser.add_argument("--end_idx", type=int, default=-1, help="End index (-1 for all)")
    parser.add_argument("--delay", type=float, default=0.05, help="Delay between API calls to avoid rate limits")
    args = parser.parse_args()

    # 创建输出目录 - 使用greedy_result目录以匹配result_postprocess.py
    # 不区分primary和plus，统一输出到标准目录
    output_dir = f"greedy_result/code_{args.task}"

    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(f"final_result/greedy_result/code_{args.task}", exist_ok=True)

    # 选择任务函数
    if args.prompt_type == "cot":
        if args.task == "debug":
            prompt_func = generate_cot_prompt_debug
        elif args.task == "translate":
            prompt_func = generate_cot_prompt_translate
        elif args.task == "polishment":
            prompt_func = generate_cot_prompt_polishment
        elif args.task == "switch":
            prompt_func = generate_cot_prompt_switch
        else:
            raise ValueError(f"Invalid task: {args.task}")
    else:
        if args.task == "debug":
            prompt_func = generate_prompt_debug
        elif args.task == "translate":
            prompt_func = generate_prompt_translate
        elif args.task == "polishment":
            prompt_func = generate_prompt_polishment
        elif args.task == "switch":
            prompt_func = generate_prompt_switch
        else:
            raise ValueError(f"Invalid task: {args.task}")

    # 初始化API处理器
    api = api_handler(args.model)

    # 读取数据 - 保留数据集类型对输入的影响，但不影响输出路径
    data_file = f"code_{args.task}"
    if args.dataset_type == "plus":
        data_file += "_plus"
    else:
        data_file += "_primary"

    data_path = f"data/{data_file}.jsonl"
    print(f"Reading data from {data_path}")

    try:
        data = read_jsonl_file(data_path)
    except FileNotFoundError:
        print(f"Error: Data file {data_path} not found!")
        print(f"Checking if base file exists...")
        fallback_path = f"data/code_{args.task}.jsonl"
        if os.path.exists(fallback_path):
            print(f"Using fallback data file: {fallback_path}")
            data = read_jsonl_file(fallback_path)
        else:
            raise FileNotFoundError(f"Neither {data_path} nor {fallback_path} could be found.")

    # 设置处理范围
    if args.end_idx == -1 or args.end_idx > len(data):
        end_idx = len(data)
    else:
        end_idx = args.end_idx
    data_subset = data[args.start_idx:end_idx]

    print(f"Processing {len(data_subset)} examples from index {args.start_idx} to {end_idx-1}")

    # 输出文件名 - 移除plus标识，使其与result_postprocess.py兼容
    model_name = args.model.replace('-', '_')
    output_file = f"{output_dir}/{model_name}.jsonl"

    print(f"Will save results to {output_file}")

    # 处理数据
    with jsonlines.open(output_file, mode='w') as writer:
        # 写入元数据作为第一行
        metadata = {
            "model": args.model,
            "task": args.task,
            "dataset_type": args.dataset_type  # 保留这个信息，但不影响输出路径
        }
        writer.write(metadata)

        # 处理每个样例
        for i, item in enumerate(tqdm(data_subset, desc=f"Processing {args.task} ({args.dataset_type}) with {args.model}")):
            try:
                # 获取提示
                if args.prompt_type == "cot":
                    prompt = prompt_func(item)
                else:
                    prompt = prompt_func(item, type=args.prompt_type)

                # 获取输出
                max_tokens = 2048
                output = api.get_output(input_text=prompt, max_tokens=max_tokens)

                # 创建结果 - 使其与vllm_inference.py格式一致
                if args.task == 'debug':
                    result = {
                        "problem_id": item.get("idx"),
                        "completion_id": 0,
                        "language": item.get("code_language"),
                        "error_type": item.get("type"),
                        "difficulty": item.get("difficulty"),
                        "prompt": prompt,
                        "code": [output],  # 保持列表格式
                    }
                elif args.task == "translate":
                    result = {
                        "problem_id": item.get("idx"),
                        "completion_id": 0,
                        "source_lang": item.get("source_lang"),
                        "target_lang": item.get("target_lang"),
                        "difficulty": item.get("difficulty"),
                        "prompt": prompt,
                        "code": [output],
                    }
                elif args.task == "polishment":
                    result = {
                        "problem_id": item.get("idx"),
                        "completion_id": 0,
                        "language": item.get("source_lang"),
                        "difficulty": item.get("difficulty"),
                        "prompt": prompt,
                        "code": [output],
                    }
                elif args.task == "switch":
                    result = {
                        "problem_id": item.get("idx"),
                        "completion_id": 0,
                        "language": item.get("language"),
                        "pair": item.get("pair_id"),
                        "prompt": prompt,
                        "code": [output],
                    }

                writer.write(result)

                # 延迟以避免API速率限制
                time.sleep(args.delay)

            except Exception as e:
                print(f"Error processing item {i + args.start_idx}: {str(e)}")

    print(f"Results saved to {output_file}")
    print(f"Total processed: {len(data_subset)} examples")
    print("You can now run result_postprocess.py to post-process the results.")

if __name__ == "__main__":
    main()