TravelPlannerLeaderboard / utils /budget_reset.py
hsaest's picture
Upload folder using huggingface_hub
9be4956 verified
import json
from evaluation.hardConstriant import get_total_cost
from tqdm import tqdm
import math
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 round_up_to_nearest_hundred(number):
return math.ceil(number / 100) * 100
if __name__ == '__main__':
query = load_line_json_data('/home/xj/toolAugEnv/code/toolConstraint/data/query/all_reset_budget.jsonl')
rewritten_query = open('/home/xj/toolAugEnv/code/toolConstraint/data/api_request/query_rewritten_reset_budget_gpt4.txt').readlines()
assert len(rewritten_query) == len(query)
for idx in tqdm(range(len(query))):
# tested_data = json.load(open(f'/home/xj/toolAugEnv/code/toolConstraint/data/annotation/all/annotation_{idx+1}.json'))
# query[idx]['reset_budget'] = round_up_to_nearest_hundred(get_total_cost(query[idx], tested_data))
# print(query[idx])
# break
query[idx]['query'] = "".join(x for x in rewritten_query[idx].strip().split('\t')[1:])
query[idx]['budget'] = query[idx]['reset_budget']
query[idx].pop('reset_budget')
with open('/home/xj/toolAugEnv/code/toolConstraint/data/query/all.jsonl', 'w') as f:
for query_unit in query:
# print(query)
json.dump(query_unit, f)
f.write('\n')
f.close()