File size: 5,928 Bytes
9be4956 |
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 |
import os
import re
import sys
sys.path.append(os.path.abspath(os.path.join(os.getcwd(), "..")))
sys.path.append('/home/xj/toolAugEnv/code/toolConstraint')
# print(sys.path)
os.chdir(os.path.dirname(os.path.abspath(__file__)))
from agents.prompts import planner_agent_prompt, cot_planner_agent_prompt, react_planner_agent_prompt,react_reflect_planner_agent_prompt,reflect_prompt
# from annotation.src.utils import get_valid_name_city,extract_before_parenthesis, extract_numbers_from_filenames
import json
import time
from langchain.callbacks import get_openai_callback
from tqdm import tqdm
from tools.planner.apis import Planner, ReactPlanner, ReactReflectPlanner
import openai
os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"
def load_line_json_data(filename):
data = []
with open(filename, 'r', encoding='utf-8') as f:
for line in f.read().strip().split('\n'):
unit = json.loads(line)
data.append(unit)
return data
def extract_numbers_from_filenames(directory):
# Define the pattern to match files
pattern = r'annotation_(\d+).json'
# List all files in the directory
files = os.listdir(directory)
# Extract numbers from filenames that match the pattern
numbers = [int(re.search(pattern, file).group(1)) for file in files if re.match(pattern, file)]
return numbers
def catch_openai_api_error():
error = sys.exc_info()[0]
if error == openai.error.APIConnectionError:
print("APIConnectionError")
elif error == openai.error.RateLimitError:
print("RateLimitError")
time.sleep(60)
elif error == openai.error.APIError:
print("APIError")
elif error == openai.error.AuthenticationError:
print("AuthenticationError")
else:
print("API error:", error)
# if __name__ == "__main__":
# user_name = 'zk'
# directory = '../../data/annotation/{}'.format(user_name)
# query_data_list = load_line_json_data('../../data/query/{}.jsonl'.format(user_name))
# numbers = extract_numbers_from_filenames(directory)
# with get_openai_callback() as cb:
# for number in tqdm(numbers[:10]):
# print(number)
# json_data = json.load(open(os.path.join(directory, 'annotation_{}.json'.format(number))))
# human_collected_info_data = json.load(open(os.path.join(directory, 'human_collected_info_{}.json'.format(number))))
# query_data = query_data_list[number-1]
# planner_results = planner.run(human_collected_info_data, query_data['query'])
# org_result = json.load(open(os.path.join('../../results/turbo16k-turbo16k/{}/plan_{}.json'.format(user_name,number))))
# # org_result.append({})
# org_result[-1]['chatgpt_human_collected_info_results'] = planner_results
# # write to json file
# # with open(os.path.join('../../results/turbo16k-turbo16k/{}/plan_{}.json'.format(user_name,number)), 'w') as f:
# # json.dump(org_result, f, indent=4)
# print(cb)
if __name__ == "__main__":
model_name=['gpt-3.5-turbo-1106','gpt-4-1106-preview','gemini','mixtral'][1]
set_type = ['dev','test'][0]
method = ['direct','cot','react','reflexion'][0]
directory = f'/home/xj/toolAugEnv/code/toolConstraint/data/final_data/{set_type}'
query_data_list = load_line_json_data(os.path.join(directory, 'query/query.jsonl'))
numbers = [i for i in range(1,len(query_data_list)+1)]
if method == 'direct':
planner = Planner(model_name=model_name, agent_prompt=planner_agent_prompt)
elif method == 'cot':
planner = Planner(model_name=model_name, agent_prompt=cot_planner_agent_prompt)
elif method == 'react':
planner = ReactPlanner(model_name=model_name, agent_prompt=react_planner_agent_prompt)
elif method == 'reflexion':
planner = ReactReflectPlanner(model_name=model_name, agent_prompt=react_reflect_planner_agent_prompt,reflect_prompt=reflect_prompt)
with get_openai_callback() as cb:
for number in tqdm(numbers[:]):
# print(number)
# json_data = json.load(open(os.path.join(directory, 'plan/annotation_{}.json'.format(number))))
human_collected_info_data = json.load(open(os.path.join(directory, 'plan/human_collected_info_{}.json'.format(number))))
query_data = query_data_list[number-1]
while True:
if method in ['react','reflexion']:
planner_results, scratchpad = planner.run(human_collected_info_data, query_data['query'])
else:
planner_results = planner.run(human_collected_info_data, query_data['query'])
if planner_results != None:
break
print(planner_results)
# check if the directory exists
if not os.path.exists(os.path.join(f'/home/xj/toolAugEnv/code/toolConstraint/results/{set_type}')):
os.makedirs(os.path.join(f'/home/xj/toolAugEnv/code/toolConstraint/results/{set_type}'))
if not os.path.exists(os.path.join(f'/home/xj/toolAugEnv/code/toolConstraint/results/{set_type}/plan_{number}.json')):
result = [{}]
else:
result = json.load(open(os.path.join(f'/home/xj/toolAugEnv/code/toolConstraint/results/{set_type}/plan_{number}.json')))
if method in ['react','reflexion']:
result[-1][f'{model_name}_{method}_collected_info_results_logs'] = scratchpad
result[-1][f'{model_name}_{method}_collected_info_results'] = planner_results
# write to json file
with open(os.path.join(f'/home/xj/toolAugEnv/code/toolConstraint/results/{set_type}/plan_{number}.json'), 'w') as f:
json.dump(result, f, indent=4)
print(cb) |