Update utils/query_element_selection.py
Browse files
utils/query_element_selection.py
CHANGED
@@ -97,7 +97,7 @@ def easy_level_element_selection(day_list):
|
|
97 |
|
98 |
|
99 |
|
100 |
-
def
|
101 |
days = random.choice(day_list)
|
102 |
date = [date.strftime('%Y-%m-%d') for date in select_consecutive_dates(days)]
|
103 |
people_number = random.choice(random.choice([[2],[3,4,5,6,7,8]]))
|
@@ -141,7 +141,7 @@ def middle_level_element_selection(day_list):
|
|
141 |
elif days == 7:
|
142 |
final_budget = round_to_hundreds(round_to_hundreds((budget["average"]+budget["highest"])/2) * people_number * 0.75)
|
143 |
|
144 |
-
query_dict = {"org": final_org, "dest": final_des, "days": days, "visiting_city_number":visiting_city_map[days], "date":date, "people_number": people_number, "local_constraint": local_constrain_record ,"budget": final_budget,"query": None, "level":"
|
145 |
return query_dict
|
146 |
|
147 |
|
@@ -212,8 +212,8 @@ def generate_elements(number:int, level="easy", day_list=[3,5,7]):
|
|
212 |
query = easy_level_element_selection(day_list)
|
213 |
if query not in query_list:
|
214 |
query_list.append(query)
|
215 |
-
elif level == "
|
216 |
-
query =
|
217 |
if query not in query_list:
|
218 |
query_list.append(query)
|
219 |
elif level == "hard":
|
@@ -229,9 +229,9 @@ def main():
|
|
229 |
|
230 |
# save query_list as jsonl file
|
231 |
for num, day_list in zip([160,160,160], [[3],[5],[7]]):
|
232 |
-
query_list = generate_elements(num,"
|
233 |
|
234 |
-
with open('../data/query/
|
235 |
for query in query_list:
|
236 |
# print(query)
|
237 |
json.dump(query, f)
|
|
|
97 |
|
98 |
|
99 |
|
100 |
+
def medium_level_element_selection(day_list):
|
101 |
days = random.choice(day_list)
|
102 |
date = [date.strftime('%Y-%m-%d') for date in select_consecutive_dates(days)]
|
103 |
people_number = random.choice(random.choice([[2],[3,4,5,6,7,8]]))
|
|
|
141 |
elif days == 7:
|
142 |
final_budget = round_to_hundreds(round_to_hundreds((budget["average"]+budget["highest"])/2) * people_number * 0.75)
|
143 |
|
144 |
+
query_dict = {"org": final_org, "dest": final_des, "days": days, "visiting_city_number":visiting_city_map[days], "date":date, "people_number": people_number, "local_constraint": local_constrain_record ,"budget": final_budget,"query": None, "level":"medium"}
|
145 |
return query_dict
|
146 |
|
147 |
|
|
|
212 |
query = easy_level_element_selection(day_list)
|
213 |
if query not in query_list:
|
214 |
query_list.append(query)
|
215 |
+
elif level == "medium":
|
216 |
+
query = medium_level_element_selection(day_list)
|
217 |
if query not in query_list:
|
218 |
query_list.append(query)
|
219 |
elif level == "hard":
|
|
|
229 |
|
230 |
# save query_list as jsonl file
|
231 |
for num, day_list in zip([160,160,160], [[3],[5],[7]]):
|
232 |
+
query_list = generate_elements(num,"medium",day_list=day_list)
|
233 |
|
234 |
+
with open('../data/query/final_annotation_medium.jsonl', 'a+') as f:
|
235 |
for query in query_list:
|
236 |
# print(query)
|
237 |
json.dump(query, f)
|