app.py

from huggingface_hub import InferenceClient
import json
import folium
from gradio_folium import Folium
import gradio as gr
from geopy.geocoders import Nominatim
from geopy.adapters import AioHTTPAdapter
import os
from urllib.parse import quote

from dotenv import load_dotenv
load_dotenv()

repo_id = "mistralai/Mistral-7B-Instruct-v0.2"
llm_client = InferenceClient(model=repo_id, timeout=180, token=os.getenv("HF_TOKEN"))

repo_id = "dslim/bert-base-NER"
NER_client = InferenceClient(model=repo_id, timeout=180, token=os.getenv("HF_TOKEN"))

def show_map(places_list):
    
    geolocator = Nominatim(user_agent="HF_smart-travel-planner") #, adapter_factory=AioHTTPAdapter)

    i = 0
    for place in places_list:
        city = place['city']
        attraction = place['attraction']
        desc = place['description']

        google_link = f"<a href='https://www.google.com/search?q={attraction}, {city}' target='_blank'><b>Learn more</b></a>"
        
        try:
            locate = geolocator.geocode(f"{attraction}, {city}", timeout=10)
        except:
            locate = None
            
        if locate is not None:
            if i == 0:
                mymap = folium.Map(location=[locate.latitude, locate.longitude], zoom_start=10)
            folium.Marker(location=[locate.latitude, locate.longitude], popup=folium.Popup(f"{attraction}: {desc} \n {google_link}", max_width=200)).add_to(mymap)
            i += 1

    return mymap


def parse_output(output):
    
    itinerary = "Itinerary" + output.split("Key points:")[0].split('Itinerary:')[1]
    key_points = output.split("Key points:")[1].strip()
    try:
        itinerary_list = json.loads(key_points)
    except:
        itinerary_list = []
        
    return itinerary, itinerary_list

def process_itinerary(itinerary, places_list):
    places = NER_client.token_classification(itinerary)
    locs = [place for place in places if place.get('entity_group') == 'LOC']
    unique_words = set()
    locs = [d for d in locs if d.get('word') not in unique_words and (unique_words.add(d.get('word')) or True)]

    words = [loc['word'] for loc in locs]
    cities = find_cities_by_attraction(places_list, words)
    
    links = [f"[{loc['word']}](https://www.google.com/search?q={quote(loc['word'])},{quote(cities[i])})" for i,loc in enumerate(locs)]

    itin_copy = itinerary

    for word, link in zip(words, links):
        itin_copy = itin_copy.replace(word, link, 1)  # Replace only the first occurrence

    return itin_copy

def find_cities_by_attraction(places, words):
    matched_cities = []
    for word in words:
        Found = False
        for place in places:
            if word.lower() in place['attraction'].lower() and not Found:  # Case insensitive match
                matched_cities.append(place['city'])
                Found = True
                
        if not Found:
            matched_cities.append(places[0]['city'].split(',')[1])
    return matched_cities
    
def run_program(place, duration, start, end):
    
    prompt = f"""Please generate a list of key cities and attractions in each cities for the following place: {place} suitable for {duration}, starting from {start} and ending at {end}, as a json list of less than 10 dictionaries with the following keys: 'city', 'attraction' and 'description'.
ALWAYS write the city and country in the 'city'. For instance do not only "city": "Paris" as the city but "city": "Paris, France". Don't combine multiplpe locations in one element. In the 'description', write the main attractions in that city as well.
Try to minimize the distance between locations. Always think of the transportation means that you want to use, and the timing: morning, afternoon, where to sleep.
Strictly generate two sections: 'Itinerary:' provides a descriptive day-wise explanation of the itinerary formatted with bullet points, then you list the locations in 'Key points:'"""

    response = llm_client.text_generation(prompt, max_new_tokens=3000) #, max_new_tokens=2000, stream=True)
    itinerary, places_list = parse_output(response)
    itinerary = process_itinerary(itinerary, places_list)
    
    mymap = show_map(places_list)

    return itinerary, mymap


with gr.Blocks(theme=gr.themes.Base(primary_hue=gr.themes.colors.green,secondary_hue=gr.themes.colors.blue)) as demo:
    
    gr.Markdown("# ✈️ Smart Travel Planner 🗺️\nThis personal travel planner is based on Mistral-7B-Instruct-v0.2, called through the Hugging Face API. Describe your destination, duration, starting and ending points of your trip. The AI assistant will suggest you an itinerary for the trip! \n Beware that the model does not have access to train or plane schedules, it is relying on general world knowledge!")
                
    dest = gr.Textbox(
        label="Destination",
        placeholder="Destination, Eg. Bavaria, Germany"
    )
    dur = gr.Textbox(
        label="Duration",
        placeholder="Duration of your trip. Eg. 3 days"
    )
    start = gr.Textbox(
        label="Starting location",
        value="anywhere",
        placeholder="Starting location of your trip. Type 'anywhere' if you don't have a place in mind"
    )
    end = gr.Textbox(
        label="Ending location",
        value="anywhere",
        placeholder="Ending location of your trip. Type 'anywhere' if you don't have a place in mind"
    )
    
    button = gr.Button("Generate itinerary!")

    gr.Markdown("### Itinerary")
    
    itinery, starting_map = run_program(dest, dur, start, end)
    
    itinerary = gr.Markdown(itinery)
    
    gr.Markdown("_Click the markers on the map to display information about the places._")
    # Get initial map
    map = Folium(value=starting_map, height=600, label="Chosen locations")

 
    # Dynamics
    button.click(run_program, inputs=[dest, dur, start, end], outputs=[itinerary, map])


if __name__ == "__main__":
    demo.launch()