110 lines
4.1 KiB
Python
110 lines
4.1 KiB
Python
import os
|
|
import gdown
|
|
import shutil
|
|
import json
|
|
from zipfile import ZipFile
|
|
|
|
urls = {
|
|
"funcqa": "https://drive.google.com/uc?id=13Sj7uIsyqWXoTh1ejWUviTzeQSES2Omd",
|
|
"restbench": "https://raw.githubusercontent.com/Yifan-Song793/RestGPT/main/datasets/tmdb.json",
|
|
"toolbench": "https://drive.google.com/uc?id=1XFjDxVZdUY7TXYF2yvzx3pJlS2fy78jk",
|
|
}
|
|
|
|
|
|
def read_jsonline(address):
|
|
not_mark = []
|
|
with open(address, 'r', encoding="utf-8") as f:
|
|
for jsonstr in f.readlines():
|
|
jsonstr = json.loads(jsonstr)
|
|
not_mark.append(jsonstr)
|
|
return not_mark
|
|
|
|
|
|
def read_json(address):
|
|
with open(address, 'r', encoding='utf-8') as json_file:
|
|
json_data = json.load(json_file)
|
|
return json_data
|
|
|
|
|
|
def toolbench_process(data_file, dataset):
|
|
ls = read_json(data_file)
|
|
all_data = read_json(f"{dataset}/tool_instruction/toolbench_tool_instruction.json")
|
|
all_dic = {}
|
|
for ID in all_data.keys():
|
|
all_dic[all_data[ID]["tool_name"]] = all_data[ID]
|
|
|
|
not_in = []
|
|
for data in ls:
|
|
Tool_dic = []
|
|
data_dic = {}
|
|
already = []
|
|
for tool in data['api_list']:
|
|
if tool['tool_name'] in all_dic:
|
|
if all_dic[tool['tool_name']]["ID"] not in already:
|
|
already.append(all_dic[tool['tool_name']]["ID"])
|
|
Tool_dic.append({"ID": all_dic[tool['tool_name']]["ID"],
|
|
"Description": all_dic[tool['tool_name']]["tool_description"], })
|
|
data["Tool_dic"] = Tool_dic
|
|
|
|
json_str = json.dumps(ls, indent=4)
|
|
with open(data_file, 'w', encoding='utf-8') as json_file:
|
|
json.dump(ls, json_file, ensure_ascii=False, indent=4)
|
|
|
|
|
|
def main():
|
|
curr_dir = os.path.dirname(__file__)
|
|
|
|
for dataset in [
|
|
"funcqa",
|
|
"restbench",
|
|
"toolbench"
|
|
]:
|
|
dataset_path = os.path.join(curr_dir, "data_{}".format(dataset), "test_data")
|
|
|
|
if not os.path.exists(dataset_path):
|
|
os.makedirs(dataset_path)
|
|
|
|
if dataset == "funcqa":
|
|
print("Processing FuncQA dataset ...\n")
|
|
# Processing FuncQA dataset ...
|
|
temp_file = os.path.join(dataset_path, "data_toolkengpt_0918.zip")
|
|
gdown.download(urls[dataset], temp_file, quiet=False)
|
|
zf = ZipFile(temp_file, 'r')
|
|
zf.extract("data/funcqa/funcqa_oh.json", ".")
|
|
zf.extract("data/funcqa/funcqa_mh.json", ".")
|
|
os.rename("data/funcqa/funcqa_oh.json", "{}/funcqa_oh.json".format(dataset_path))
|
|
os.rename("data/funcqa/funcqa_mh.json", "{}/funcqa_mh.json".format(dataset_path))
|
|
os.remove(temp_file)
|
|
shutil.rmtree("data")
|
|
|
|
print("FuncQA dataset Done!\n")
|
|
|
|
if dataset == "restbench":
|
|
print("Processing RestBench dataset ... \n")
|
|
# Processing RestBench Dataset
|
|
os.system("wget -P {} -c {}".format(dataset_path, urls[dataset]))
|
|
|
|
print("RestBench dataset Done!\n")
|
|
|
|
if dataset == "toolbench":
|
|
print("Processing ToolBench dataset ... \n")
|
|
# Processing ToolBench Dataset
|
|
temp_file = os.path.join(dataset_path, "data.zip")
|
|
gdown.download(urls[dataset], temp_file, quiet=False)
|
|
zf = ZipFile(temp_file, 'r')
|
|
zf.extract("data/test_instruction/G2_category.json", ".")
|
|
zf.extract("data/test_instruction/G3_instruction.json", ".")
|
|
|
|
os.rename("data/test_instruction/G2_category.json", "{}/G2_category.json".format(dataset_path))
|
|
os.rename("data/test_instruction/G3_instruction.json", "{}/G3_instruction.json".format(dataset_path))
|
|
toolbench_process("{}/G2_category.json".format(dataset_path), "data_{}".format(dataset))
|
|
toolbench_process("{}/G3_instruction.json".format(dataset_path), "data_{}".format(dataset))
|
|
os.remove(temp_file)
|
|
shutil.rmtree("data")
|
|
|
|
print("Toolbench dataset Done!\n")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|