wehub-resource-sync
91 lines
3.7 KiB
Python
91 lines
3.7 KiB
Python
# — coding: utf-8 –
|
||
import openai
|
||
import json
|
||
import argparse
|
||
import os
|
||
from tqdm import tqdm
|
||
from easytool import funcQA, restbench, toolbench_retrieve, toolbench
|
||
from easytool.util import *
|
||
openai.api_key = os.environ["OPENAI_API_KEY"]
|
||
|
||
if __name__ == '__main__':
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--model_name', type=str, default='gpt-3.5-turbo')
|
||
parser.add_argument('--task', type=str, default='funcqa_mh', help='funcqa, toolbench_retrieve, toolbench, restbench')
|
||
parser.add_argument('--data_type', type=str, default='G3', help='G2 or G3 or funcqa_mh or funcqa_oh')
|
||
parser.add_argument('--tool_root_dir', type=str, default='.toolenv/tools/')
|
||
parser.add_argument('--retrieval_num', type=int, default=5)
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.task == 'funcqa':
|
||
dataset = read_json('data_funcqa/tool_instruction/functions_data.json')
|
||
Tool_dic = read_jsonline('data_funcqa/tool_instruction/tool_dic.jsonl')
|
||
test_data = read_json(f"data_funcqa/test_data/{args.data_type}.json")
|
||
progress_file = f"FuncQA_{args.data_type}_{args.model_name}_Easytool.txt"
|
||
|
||
|
||
elif 'toolbench' in args.task:
|
||
base_path = args.tool_root_dir
|
||
index = build_index(base_path)
|
||
dataset = read_json('data_toolbench/tool_instruction/toolbench_tool_instruction.json')
|
||
if args.data_type == 'G2':
|
||
test_data = read_json(f'''data_toolbench/test_data/{args.data_type}_category.json''')
|
||
elif args.data_type == 'G3':
|
||
test_data = read_json(f'''data_toolbench/test_data/{args.data_type}_instruction.json''')
|
||
progress_file = f'''{args.data_type}_{args.model_name}_Easytool.txt'''
|
||
|
||
|
||
elif args.task == 'restbench':
|
||
Tool_dic = read_json('data_restbench/tool_instruction/tmdb_tool.json')
|
||
dic_tool = {}
|
||
for data in Tool_dic:
|
||
dic_tool[data['ID']] = data
|
||
test_data = read_json('data_restbench/test_data/tmdb.json')
|
||
progress_file = f"restbench_{args.model_name}_Easytool.txt"
|
||
|
||
else:
|
||
print("Wrong task name")
|
||
exit()
|
||
|
||
start_index = get_last_processed_index(progress_file)
|
||
total_files = len(test_data)
|
||
retrieval_num = args.retrieval_num
|
||
ind = start_index
|
||
model_name = args.model_name
|
||
|
||
print("-------Start Execution-------")
|
||
if args.data_type == 'funcqa_mh':
|
||
funcQA.task_execution_mh(args.data_type, start_index, total_files,
|
||
retrieval_num, ind, model_name, dataset,
|
||
Tool_dic, test_data, progress_file)
|
||
elif args.data_type == 'funcqa_oh':
|
||
funcQA.task_execution_oh(args.data_type, start_index, total_files,
|
||
retrieval_num, ind, model_name, dataset,
|
||
Tool_dic, test_data, progress_file)
|
||
|
||
|
||
elif args.task == 'toolbench_retrieve':
|
||
toolbench_retrieve.task_execution(args.data_type,
|
||
base_path, index, dataset, test_data, progress_file,
|
||
start_index, total_files, retrieval_num, ind, model_name)
|
||
|
||
|
||
|
||
elif args.task == 'toolbench':
|
||
toolbench.task_execution(args.data_type,
|
||
base_path, index, dataset, test_data, progress_file,
|
||
start_index, total_files, retrieval_num, ind, model_name)
|
||
|
||
|
||
|
||
elif args.task == 'restbench':
|
||
restbench.task_execution(
|
||
Tool_dic, dic_tool, test_data, progress_file,
|
||
start_index, total_files, retrieval_num, ind, model_name)
|
||
|
||
|
||
else:
|
||
print("Wrong task name")
|
||
exit()
|