from pathlib import Path from json import load as load_json from numpy.random import permutation as perm data_dir = Path(__file__).parent / "data" / "cmu" / "processed" def get_cities(): with open(data_dir / "cities.json", "r") as f: return load_json(f) def get_score_threshold_per_city(): with open(data_dir / "score_threshold_per_city.json", "r") as f: return load_json(f) def get_city_to_hotel_id_map(): with open(data_dir / "city_to_hotel_id_map.json", "r") as f: return load_json(f) def get_hotel_id_to_name_map(): with open(data_dir / "hotel_id_to_name_map.json", "r") as f: return load_json(f) def get_hotel_id_to_review_map(): with open(data_dir / "hotel_id_to_review_map.json", "r") as f: return load_json(f) score_threshold_per_city = get_score_threshold_per_city() city_to_hotel_id_map = get_city_to_hotel_id_map() hotel_id_to_name_map = get_hotel_id_to_name_map() hotel_id_to_review_map = get_hotel_id_to_review_map() def get_reviews_for_prompt(city, preferences) -> dict: for hotel_id in perm(city_to_hotel_id_map[city]): hotel_id = str(hotel_id) res = {"hotel_name": hotel_id_to_name_map[hotel_id], 'positive': [], 'negative': []} try: hotel_reviews = hotel_id_to_review_map[hotel_id]['reviews'] except KeyError: continue for review in perm(hotel_reviews): if (review['score'] == 5) & (len(res['positive']) < 3): res['positive'].append(review) if (review['score'] <= 2) & (len(res['negative']) < 1): res['negative'].append(review) if (len(res['positive']) >= 3) & (len(res['negative']) >= 1): return res return None