utils/extract_human_annotation.py

import json
import os
import re
import sys
from tools.flights.apis import Flights
from tools.accommodations.apis import Accommodations
from tools.restaurants.apis import Restaurants
from tools.googleDistanceMatrix.apis import GoogleDistanceMatrix
from tools.googlePlaces.apis import GooglePlaces
from tools.attractions.apis import Attractions
from annotation.src.utils import get_valid_name_city,extract_before_parenthesis
from tqdm import tqdm

sys.path.append(os.path.abspath(os.path.join(os.getcwd(), "..")))
os.chdir(os.path.dirname(os.path.abspath(__file__)))


flight = Flights()
accommodations = Accommodations()
restaurants = Restaurants()
googleDistanceMatrix = GoogleDistanceMatrix()
googlePlaces = GooglePlaces()
attractions = Attractions()


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 extract_from_to(text: str):
    """
    Extracts 'A' and 'B' from the format "from A to B" in the given text, with B ending at a comma or the end of the string.
    
    Args:
    - text (str): The input string.
    
    Returns:
    - tuple: A tuple containing 'A' and 'B'. If no match is found, returns (None, None).
    """
    pattern = r"from\s+(.+?)\s+to\s+([^,]+)(?=[,\s]|$)"
    matches = re.search(pattern, text)
    return matches.groups() if matches else (None, None)

def extract_city_list(query_data, annotated_data):
    city_list = []
    for unit in annotated_data[:query_data['days']]:
        if 'from' in unit['current_city']:
            from_city, to_city = extract_from_to(unit['current_city'])
            from_city = extract_before_parenthesis(from_city)
            to_city = extract_before_parenthesis(to_city)
            if from_city not in city_list:
                city_list.append(from_city)
            if to_city not in city_list:
                city_list.append(to_city)
        else:
            city = extract_before_parenthesis(unit['current_city'])
            if city not in city_list:
                city_list.append(city)

    return city_list


# if __name__ == '__main__':
#     user_name = 'all'
#     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)
#     print(numbers)
#     for number in tqdm(numbers):
#         json_data = json.load(open(os.path.join(directory, 'annotation_{}.json'.format(number))))
#         query_data = query_data_list[number-1]
#         city_list = extract_city_list(query_data,json_data)
#         human_collected_info = []

#         for city in city_list[1:]:
#             attractions_data = attractions.run(city)
#             if type(attractions_data) != str:
#                 attractions_data.drop(['Latitude','Longitude','Address','Phone','Website','City'],axis=1,inplace=True)
#             if type(attractions_data) != str:
#                 attractions_data = attractions_data.to_string(index=False)
#             restaurants_data = restaurants.run(city)
#             restaurants_data.drop(['City'],axis=1,inplace=True)
#             if type(restaurants_data) != str:
#                 restaurants_data = restaurants_data.to_string(index=False)
#             accommodations_data = accommodations.run(city)
#             accommodations_data.drop(['city'],axis=1,inplace=True)
#             if type(accommodations_data) != str:
#                 accommodations_data = accommodations_data.to_string(index=False)
#             human_collected_info.append({"Description":"Attractions in {}".format(city),"Content":attractions_data})
#             human_collected_info.append({"Description":"Restaurants in {}".format(city),"Content":restaurants_data})
#             human_collected_info.append({"Description":"Accommodations in {}".format(city),"Content":accommodations_data})

        
#         for idx, unit in enumerate(json_data):
#             if unit != {}:
#                 if 'from' in unit['current_city']:
#                     from_city, to_city = extract_from_to(unit['current_city'])
#                     from_city = extract_before_parenthesis(from_city)
#                     to_city = extract_before_parenthesis(to_city)
#                     date = query_data_list[number-1]['date'][idx]
#                     flight_data = flight.run(from_city, to_city, date)
#                     if type(flight_data) != str:
#                         flight_data.drop(['OriginCityName','DestCityName','Distance','FlightDate'],axis=1,inplace=True)
#                         flight_data = flight_data.to_string(index=False)
#                     human_collected_info.append({"Description":"Flight from {} to {} on {}".format(from_city, to_city, date), "Content":flight_data})
#                     self_driving_data = googleDistanceMatrix.run(from_city, to_city,mode="self-driving")
#                     human_collected_info.append({"Description":"Self-driving from {} to {}".format(from_city, to_city), "Content":self_driving_data})
#                     taxi_data = googleDistanceMatrix.run(from_city, to_city, mode='taxi')
#                     human_collected_info.append({"Description":"Taxi from {} to {}".format(from_city, to_city), "Content":taxi_data})
        
#         # write to json file
#         with open(os.path.join(directory, 'human_collected_info_{}.json'.format(number)), 'w', encoding='utf-8') as f:
#             json.dump(human_collected_info, f, indent=4, ensure_ascii=False)
#         # break


if __name__ == '__main__':
    set_type = ['train','dev','test'][2]
    directory = '/home/xj/toolAugEnv/code/toolConstraint/data/final_data/{}'.format(set_type)
    query_data_list = load_line_json_data('/home/xj/toolAugEnv/code/toolConstraint/data/final_data/{}/query/query.jsonl'.format(set_type))
    numbers = [i for i in range(1,len(query_data_list)+1)]
    for number in tqdm(numbers):
        json_data = json.load(open(os.path.join(directory, 'plan/plan_{}.json'.format(number))))[1]
        query_data = query_data_list[number-1]
        city_list = extract_city_list(query_data,json_data)
        human_collected_info = []

        for city in city_list[1:]:
            attractions_data = attractions.run(city)
            # if type(attractions_data) != str:
            #     attractions_data.drop(['Latitude','Longitude','Address','Phone','Website','City'],axis=1,inplace=True)
            if type(attractions_data) != str:
                attractions_data = attractions_data.to_string(index=False)
            restaurants_data = restaurants.run(city)
            # restaurants_data.drop(['City'],axis=1,inplace=True)
            if type(restaurants_data) != str:
                restaurants_data = restaurants_data.to_string(index=False)
            accommodations_data = accommodations.run(city)
            # accommodations_data.drop(['city'],axis=1,inplace=True)
            if type(accommodations_data) != str:
                accommodations_data = accommodations_data.to_string(index=False)
            human_collected_info.append({"Description":"Attractions in {}".format(city),"Content":attractions_data})
            human_collected_info.append({"Description":"Restaurants in {}".format(city),"Content":restaurants_data})
            human_collected_info.append({"Description":"Accommodations in {}".format(city),"Content":accommodations_data})

        
        for idx, unit in enumerate(json_data):
            if unit != {}:
                if 'from' in unit['current_city']:
                    from_city, to_city = extract_from_to(unit['current_city'])
                    from_city = extract_before_parenthesis(from_city)
                    to_city = extract_before_parenthesis(to_city)
                    date = query_data_list[number-1]['date'][idx]
                    flight_data = flight.run(from_city, to_city, date)
                    if type(flight_data) != str:
                        # flight_data.drop(['OriginCityName','DestCityName','Distance','FlightDate'],axis=1,inplace=True)
                        flight_data = flight_data.to_string(index=False)
                    human_collected_info.append({"Description":"Flight from {} to {} on {}".format(from_city, to_city, date), "Content":flight_data})
                    self_driving_data = googleDistanceMatrix.run(from_city, to_city,mode="self-driving")
                    human_collected_info.append({"Description":"Self-driving from {} to {}".format(from_city, to_city), "Content":self_driving_data})
                    taxi_data = googleDistanceMatrix.run(from_city, to_city, mode='taxi')
                    human_collected_info.append({"Description":"Taxi from {} to {}".format(from_city, to_city), "Content":taxi_data})
        
        # write to json file
        with open(os.path.join(directory, 'plan/human_collected_info_{}.json'.format(number)), 'w', encoding='utf-8') as f:
            json.dump(human_collected_info, f, indent=4, ensure_ascii=False)
        # break