microsoft--jarvis
wehub-resource-sync
91 行
3.7 KiB
Python
91 行
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()
|